llama_cpp_4/sampling.rs
1//! Safe wrapper around `llama_sampler`.
2
3use std::borrow::Borrow;
4use std::ffi::{c_char, CString};
5use std::fmt::{Debug, Formatter};
6use std::ptr::NonNull;
7
8use llama_cpp_sys_4::{
9 common::common_sampler_params, llama_logit_bias, llama_sampler,
10 llama_sampler_chain_add, llama_sampler_chain_default_params, llama_sampler_chain_init,
11 llama_sampler_chain_n, llama_sampler_chain_remove, llama_sampler_clone, llama_sampler_copy,
12 llama_sampler_free, llama_sampler_get_seed, llama_sampler_init_adaptive_p,
13 llama_sampler_init_dist, llama_sampler_init_dry, llama_sampler_init_grammar,
14 llama_sampler_init_grammar_lazy_patterns, llama_sampler_init_greedy, llama_sampler_init_infill,
15 llama_sampler_init_logit_bias, llama_sampler_init_min_p, llama_sampler_init_mirostat,
16 llama_sampler_init_mirostat_v2, llama_sampler_init_penalties, llama_sampler_init_temp,
17 llama_sampler_init_temp_ext, llama_sampler_init_top_k, llama_sampler_init_top_n_sigma,
18 llama_sampler_init_top_p, llama_sampler_init_typical, llama_sampler_init_xtc,
19 llama_sampler_name, llama_sampler_reset,
20};
21
22use crate::context::LlamaContext;
23use crate::model::LlamaModel;
24use crate::token::data_array::LlamaTokenDataArray;
25use crate::token::LlamaToken;
26
27/// Why a sampler could not be constructed.
28///
29/// The panicking constructors — [`LlamaSampler::grammar`] and friends — are
30/// kept for the common case where the grammar is a literal or came from
31/// [`ChatParams::grammar`](crate::chat::ChatParams), and a failure really is a
32/// bug. These exist for the case they handle badly: a grammar, rule name or
33/// trigger pattern that arrived from a request, where an interior NUL is bad
34/// *input* rather than a broken invariant.
35#[derive(Debug, PartialEq, Eq, thiserror::Error)]
36pub enum SamplerInitError {
37 /// A caller-supplied string contained an interior NUL, so it cannot cross
38 /// the FFI boundary. `field` names which argument.
39 #[error("{field} contains an interior NUL byte")]
40 InteriorNul {
41 /// Name of the offending argument.
42 field: &'static str,
43 /// The underlying conversion error.
44 #[source]
45 source: std::ffi::NulError,
46 },
47 /// llama.cpp returned null. For a grammar sampler this is what an
48 /// unparseable GBNF grammar looks like; it logs the parse error itself.
49 #[error("llama.cpp could not build the {0} sampler (see the llama.cpp log)")]
50 NullReturn(&'static str),
51}
52
53/// Convert a caller string, naming the argument so the error says which one.
54fn c_string(field: &'static str, value: &str) -> Result<CString, SamplerInitError> {
55 CString::new(value).map_err(|source| SamplerInitError::InteriorNul { field, source })
56}
57
58/// A safe wrapper around `llama_sampler`.
59pub struct LlamaSampler {
60 pub(crate) sampler: NonNull<llama_sampler>,
61}
62
63impl Debug for LlamaSampler {
64 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
65 f.debug_struct("LlamaSamplerChain").finish()
66 }
67}
68#[derive(Debug, Clone)]
69#[allow(
70 missing_docs,
71 clippy::struct_excessive_bools,
72 clippy::module_name_repetitions,
73 dead_code
74)]
75pub struct LlamaSamplerParams {
76 top_k: i32,
77 top_p: f32,
78 temp: f32,
79 seed: u32,
80}
81
82impl LlamaSamplerParams {
83 /// Set the seed of the context
84 ///
85 /// # Examples
86 ///
87 /// ```rust
88 /// use llama_cpp_4::sampling::LlamaSamplerParams;
89 /// let params = LlamaSamplerParams::default();
90 /// let params = params.with_seed(1234);
91 /// assert_eq!(params.seed(), 1234);
92 /// ```
93 #[must_use]
94 pub fn with_seed(mut self, seed: u32) -> Self {
95 self.seed = seed;
96 self
97 }
98
99 /// Get the seed of the context
100 ///
101 /// # Examples
102 ///
103 /// ```rust
104 /// use llama_cpp_4::sampling::LlamaSamplerParams;
105 /// let params = LlamaSamplerParams::default()
106 /// .with_seed(1234);
107 /// assert_eq!(params.seed(), 1234);
108 /// ```
109 #[must_use]
110 pub fn seed(&self) -> u32 {
111 self.seed
112 }
113}
114
115impl Default for LlamaSamplerParams {
116 fn default() -> Self {
117 Self {
118 top_k: 50,
119 top_p: 0.9,
120 temp: 0.8,
121 seed: 1234,
122 }
123 }
124}
125
126impl Default for LlamaSampler {
127 fn default() -> Self {
128 Self::new()
129 }
130}
131
132impl LlamaSampler {
133 /// Create new sampler with default params.
134 ///
135 /// # Panics
136 ///
137 /// Panics if llama.cpp returns a null pointer.
138 #[must_use]
139 pub fn new() -> Self {
140 let sparams = unsafe { llama_sampler_chain_default_params() };
141
142 Self {
143 sampler: NonNull::new(unsafe { llama_sampler_chain_init(sparams) }).unwrap(),
144 }
145 }
146
147 /// Sample and accept a token from the idx-th output of the last evaluation
148 #[must_use]
149 /// # Panics
150 ///
151 /// Panics if llama.cpp's grammar code rejects the state — see
152 /// [`Self::try_sample`], which returns the error instead. This *panics*
153 /// rather than aborting because the call is routed through a guard; calling
154 /// `llama_sampler_sample` directly would let a C++ exception unwind into
155 /// Rust and kill the process.
156 pub fn sample(&self, ctx: &LlamaContext, idx: i32) -> LlamaToken {
157 self.try_sample(ctx, idx)
158 .unwrap_or_else(|e| panic!("sampling failed: {e}"))
159 }
160
161 /// Sample a token, reporting llama.cpp's failures instead of panicking.
162 ///
163 /// The realistic failure is a grammar that can no longer accept anything:
164 /// llama.cpp raises *"Unexpected empty grammar stack"* when a model's
165 /// vocabulary cannot satisfy the constraint — a JSON schema against a
166 /// vocabulary with no `{`, say. That is a property of the model/grammar
167 /// pair, not a bug, so it is worth handling.
168 ///
169 /// # Errors
170 ///
171 /// Returns [`ShimError::Failed`](crate::shim::ShimError::Failed) carrying
172 /// llama.cpp's message.
173 pub fn try_sample(
174 &self,
175 ctx: &LlamaContext,
176 idx: i32,
177 ) -> Result<LlamaToken, crate::shim::ShimError> {
178 let mut status = llama_cpp_sys_4::LLAMA_SHIM_OK;
179 let token = unsafe {
180 llama_cpp_sys_4::common_shim_sampler_sample_raw(
181 self.sampler.as_ptr(),
182 ctx.context.as_ptr(),
183 idx,
184 &raw mut status,
185 )
186 };
187 crate::shim::check_status(status)?;
188 Ok(LlamaToken(token))
189 }
190
191 /// Applies this sampler to a [`LlamaTokenDataArray`].
192 pub fn apply(&mut self, data_array: &mut LlamaTokenDataArray) {
193 data_array.apply_sampler(self);
194 }
195
196 /// Accepts a token from the sampler, possibly updating the internal state of certain samplers
197 /// (e.g. grammar, repetition, etc.)
198 /// # Panics
199 ///
200 /// Panics if the token leaves a grammar with nothing it can accept — see
201 /// [`Self::try_accept`].
202 pub fn accept(&mut self, token: LlamaToken) {
203 self.try_accept(token)
204 .unwrap_or_else(|e| panic!("accepting token {}: {e}", token.0));
205 }
206
207 /// Accept a token, reporting llama.cpp's failures instead of panicking.
208 ///
209 /// # Errors
210 ///
211 /// Returns [`ShimError::Failed`](crate::shim::ShimError::Failed) when the
212 /// grammar cannot accept this token — see [`Self::try_sample`].
213 pub fn try_accept(&mut self, token: LlamaToken) -> Result<(), crate::shim::ShimError> {
214 let status = unsafe {
215 llama_cpp_sys_4::common_shim_sampler_accept_raw(self.sampler.as_ptr(), token.0)
216 };
217 crate::shim::check_status(status)
218 }
219
220 /// Accepts several tokens from the sampler or context, possibly updating the internal state of
221 /// certain samplers (e.g. grammar, repetition, etc.)
222 /// # Panics
223 ///
224 /// Panics on the first token the grammar cannot accept — see
225 /// [`Self::try_accept`].
226 pub fn accept_many(&mut self, tokens: impl IntoIterator<Item = impl Borrow<LlamaToken>>) {
227 for token in tokens {
228 self.accept(*token.borrow());
229 }
230 }
231
232 /// Accepts several tokens from the sampler or context, possibly updating the internal state of
233 /// certain samplers (e.g. grammar, repetition, etc.)
234 #[must_use]
235 pub fn with_tokens(
236 mut self,
237 tokens: impl IntoIterator<Item = impl Borrow<LlamaToken>>,
238 ) -> Self {
239 self.accept_many(tokens);
240 self
241 }
242
243 /// Combines a list of samplers into a single sampler that applies each component sampler one
244 /// after another.
245 ///
246 /// If you are using a chain to select a token, the chain should always end with one of
247 /// [`LlamaSampler::greedy`], [`LlamaSampler::dist`], [`LlamaSampler::mirostat`], and
248 /// [`LlamaSampler::mirostat_v2`].
249 ///
250 /// # Panics
251 ///
252 /// Panics if llama.cpp returns a null pointer.
253 #[must_use]
254 pub fn chain(samplers: impl IntoIterator<Item = Self>, no_perf: bool) -> Self {
255 unsafe {
256 let mut params = llama_sampler_chain_default_params();
257 params.no_perf = no_perf;
258 let chain = llama_sampler_chain_init(params);
259
260 for sampler in samplers {
261 llama_sampler_chain_add(chain, sampler.sampler.as_ptr());
262
263 // Do not call `llama_sampler_free` on the sampler, as the internal sampler is now
264 // owned by the chain
265 std::mem::forget(sampler);
266 }
267
268 Self {
269 sampler: NonNull::new(chain).unwrap(),
270 }
271 }
272 }
273
274 /// Same as [`Self::chain`] with `no_perf = false`.
275 ///
276 /// # Panics
277 ///
278 /// Panics if llama.cpp returns a null pointer.
279 ///
280 /// # Example
281 /// ```rust
282 /// use llama_cpp_4::token::{
283 /// LlamaToken,
284 /// data::LlamaTokenData,
285 /// data_array::LlamaTokenDataArray
286 /// };
287 /// use llama_cpp_4::sampling::LlamaSampler;
288 ///
289 /// let mut data_array = LlamaTokenDataArray::new(vec![
290 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
291 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
292 /// LlamaTokenData::new(LlamaToken(2), 2., 0.),
293 /// ], false);
294 ///
295 /// data_array.apply_sampler(&mut LlamaSampler::chain_simple([
296 /// LlamaSampler::temp(0.5),
297 /// LlamaSampler::greedy(),
298 /// ]));
299 ///
300 /// assert_eq!(data_array.data[0].logit(), 0.);
301 /// assert_eq!(data_array.data[1].logit(), 2.);
302 /// assert_eq!(data_array.data[2].logit(), 4.);
303 ///
304 /// assert_eq!(data_array.data.len(), 3);
305 /// assert_eq!(data_array.selected_token(), Some(LlamaToken(2)));
306 /// ```
307 #[must_use]
308 pub fn chain_simple(samplers: impl IntoIterator<Item = Self>) -> Self {
309 Self::chain(samplers, false)
310 }
311
312 /// Updates the logits `l_i`' = `l_i/t`. When `t <= 0.0`, the maximum logit is kept at its original
313 /// value, the rest are set to -inf.
314 ///
315 /// # Panics
316 ///
317 /// Panics if llama.cpp returns a null pointer.
318 ///
319 /// # Example:
320 /// ```rust
321 /// use llama_cpp_4::token::{
322 /// LlamaToken,
323 /// data::LlamaTokenData,
324 /// data_array::LlamaTokenDataArray
325 /// };
326 /// use llama_cpp_4::sampling::LlamaSampler;
327 ///
328 /// let mut data_array = LlamaTokenDataArray::new(vec![
329 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
330 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
331 /// LlamaTokenData::new(LlamaToken(2), 2., 0.),
332 /// ], false);
333 ///
334 /// data_array.apply_sampler(&mut LlamaSampler::temp(0.5));
335 ///
336 /// assert_eq!(data_array.data[0].logit(), 0.);
337 /// assert_eq!(data_array.data[1].logit(), 2.);
338 /// assert_eq!(data_array.data[2].logit(), 4.);
339 /// ```
340 #[must_use]
341 pub fn temp(t: f32) -> Self {
342 let sampler = unsafe { llama_sampler_init_temp(t) };
343 Self {
344 sampler: NonNull::new(sampler).unwrap(),
345 }
346 }
347
348 /// Dynamic temperature implementation (a.k.a. entropy) described in the paper
349 /// <https://arxiv.org/abs/2309.02772>.
350 ///
351 /// # Panics
352 ///
353 /// Panics if llama.cpp returns a null pointer.
354 #[must_use]
355 pub fn temp_ext(t: f32, delta: f32, exponent: f32) -> Self {
356 let sampler = unsafe { llama_sampler_init_temp_ext(t, delta, exponent) };
357 Self {
358 sampler: NonNull::new(sampler).unwrap(),
359 }
360 }
361
362 /// Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration"
363 /// <https://arxiv.org/abs/1904.09751>.
364 ///
365 /// # Panics
366 ///
367 /// Panics if llama.cpp returns a null pointer.
368 ///
369 /// # Example:
370 /// ```rust
371 /// use llama_cpp_4::token::{
372 /// LlamaToken,
373 /// data::LlamaTokenData,
374 /// data_array::LlamaTokenDataArray
375 /// };
376 /// use llama_cpp_4::sampling::LlamaSampler;
377 ///
378 /// let mut data_array = LlamaTokenDataArray::new(vec![
379 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
380 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
381 /// LlamaTokenData::new(LlamaToken(2), 2., 0.),
382 /// LlamaTokenData::new(LlamaToken(3), 3., 0.),
383 /// ], false);
384 ///
385 /// data_array.apply_sampler(&mut LlamaSampler::top_k(2));
386 ///
387 /// assert_eq!(data_array.data.len(), 2);
388 /// assert_eq!(data_array.data[0].id(), LlamaToken(3));
389 /// assert_eq!(data_array.data[1].id(), LlamaToken(2));
390 /// ```
391 #[must_use]
392 pub fn top_k(k: i32) -> Self {
393 let sampler = unsafe { llama_sampler_init_top_k(k) };
394 Self {
395 sampler: NonNull::new(sampler).unwrap(),
396 }
397 }
398
399 /// Locally Typical Sampling implementation described in the paper <https://arxiv.org/abs/2202.00666>.
400 ///
401 /// # Panics
402 ///
403 /// Panics if llama.cpp returns a null pointer.
404 #[must_use]
405 pub fn typical(p: f32, min_keep: usize) -> Self {
406 let sampler = unsafe { llama_sampler_init_typical(p, min_keep) };
407 Self {
408 sampler: NonNull::new(sampler).unwrap(),
409 }
410 }
411
412 /// Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration"
413 /// <https://arxiv.org/abs/1904.09751>.
414 ///
415 /// # Panics
416 ///
417 /// Panics if llama.cpp returns a null pointer.
418 #[must_use]
419 pub fn top_p(p: f32, min_keep: usize) -> Self {
420 let sampler = unsafe { llama_sampler_init_top_p(p, min_keep) };
421 Self {
422 sampler: NonNull::new(sampler).unwrap(),
423 }
424 }
425
426 /// Minimum P sampling as described in <https://github.com/ggerganov/llama.cpp/pull/3841>.
427 ///
428 /// # Panics
429 ///
430 /// Panics if llama.cpp returns a null pointer.
431 #[must_use]
432 pub fn min_p(p: f32, min_keep: usize) -> Self {
433 let sampler = unsafe { llama_sampler_init_min_p(p, min_keep) };
434 Self {
435 sampler: NonNull::new(sampler).unwrap(),
436 }
437 }
438
439 /// XTC sampler as described in <https://github.com/oobabooga/text-generation-webui/pull/6335>.
440 ///
441 /// # Panics
442 ///
443 /// Panics if llama.cpp returns a null pointer.
444 #[must_use]
445 pub fn xtc(p: f32, t: f32, min_keep: usize, seed: u32) -> Self {
446 let sampler = unsafe { llama_sampler_init_xtc(p, t, min_keep, seed) };
447 Self {
448 sampler: NonNull::new(sampler).unwrap(),
449 }
450 }
451
452 /// Grammar sampler
453 ///
454 /// # Panics
455 /// - If either of `grammar_str` or `grammar_root` contain null bytes.
456 /// - If llama.cpp returns a null pointer, which is also what an
457 /// unparseable GBNF grammar looks like.
458 ///
459 /// Prefer [`Self::try_grammar`] when the grammar did not come from your own
460 /// source — a grammar built from a request is input, and both failures above
461 /// are then reachable without a bug.
462 #[must_use]
463 pub fn grammar(model: &LlamaModel, grammar_str: &str, grammar_root: &str) -> Self {
464 Self::try_grammar(model, grammar_str, grammar_root)
465 .unwrap_or_else(|e| panic!("building grammar sampler: {e}"))
466 }
467
468 /// [`Self::grammar`], reporting llama.cpp's and the conversion's failures
469 /// instead of panicking.
470 ///
471 /// # Errors
472 ///
473 /// [`SamplerInitError::InteriorNul`] if either string contains a NUL, or
474 /// [`SamplerInitError::NullReturn`] if llama.cpp rejected the grammar.
475 pub fn try_grammar(
476 model: &LlamaModel,
477 grammar_str: &str,
478 grammar_root: &str,
479 ) -> Result<Self, SamplerInitError> {
480 let grammar_str = c_string("grammar_str", grammar_str)?;
481 let grammar_root = c_string("grammar_root", grammar_root)?;
482
483 // SAFETY: both pointers come from live `CString`s that outlast the call,
484 // and llama.cpp only reads them. The vocab reference is borrowed from
485 // `model` for the duration.
486 let sampler = unsafe {
487 llama_sampler_init_grammar(
488 model.get_vocab().vocab.as_ref(),
489 grammar_str.as_ptr(),
490 grammar_root.as_ptr(),
491 )
492 };
493 NonNull::new(sampler)
494 .map(|sampler| Self { sampler })
495 .ok_or(SamplerInitError::NullReturn("grammar"))
496 }
497
498 /// DRY sampler, designed by p-e-w, as described in:
499 /// <https://github.com/oobabooga/text-generation-webui/pull/5677>, porting Koboldcpp
500 /// implementation authored by pi6am: <https://github.com/LostRuins/koboldcpp/pull/982>
501 ///
502 /// # Panics
503 /// - If any string in `seq_breakers` contains null bytes.
504 /// - If llama.cpp returns a null pointer.
505 #[allow(clippy::too_many_arguments)]
506 #[must_use]
507 pub fn dry(
508 &self,
509 model: &LlamaModel,
510 multiplier: f32,
511 base: f32,
512 allowed_length: i32,
513 penalty_last_n: i32,
514 seq_breakers: impl IntoIterator<Item = impl AsRef<[u8]>>,
515 ) -> Self {
516 // A NUL in a sequence breaker is the same class of input error the
517 // grammar constructors have a `try_` variant for; `dry` keeps the panic
518 // because its breakers are conventionally literals (`"\n"`, `":"`).
519 let seq_breakers: Vec<CString> = seq_breakers
520 .into_iter()
521 .map(|s| CString::new(s.as_ref()).expect("sequence breaker contains an interior NUL"))
522 .collect();
523 // CString::as_ptr() returns *const c_char, which matches what the binding
524 // expects on every platform (signed on macOS/x86 Linux, unsigned on musl ARM).
525 let mut seq_breaker_pointers: Vec<*const c_char> =
526 seq_breakers.iter().map(|s| s.as_ptr()).collect();
527
528 let sampler = unsafe {
529 llama_sampler_init_dry(
530 model.get_vocab().vocab.as_ref(),
531 multiplier,
532 base,
533 allowed_length,
534 penalty_last_n,
535 seq_breaker_pointers.as_mut_ptr(),
536 seq_breaker_pointers.len(),
537 )
538 };
539
540 Self {
541 sampler: NonNull::new(sampler).unwrap(),
542 }
543 }
544
545 /// Penalizes tokens for being present in the context.
546 ///
547 /// Parameters:
548 /// - `n_vocab`: [`LlamaModel::n_vocab`]
549 /// - `penalty_last_n`: last n tokens to penalize (0 = disable penalty)
550 /// - `penalty_repeat`: repetition penalty (must be > 0.0, 1.0 = disabled)
551 /// - `penalty_freq`: frequency penalty (must be finite, 0.0 = disabled)
552 /// - `penalty_present`: presence penalty (must be finite, 0.0 = disabled)
553 ///
554 /// If `penalty_last_n` is `0`, or every penalty sits at its disabled value,
555 /// llama.cpp returns a no-op sampler named `"?penalties"` — see
556 /// [`Self::name`].
557 ///
558 /// # Panics
559 ///
560 /// Panics if llama.cpp returns a null pointer.
561 #[allow(clippy::too_many_arguments)]
562 #[must_use]
563 pub fn penalties(
564 n_vocab: i32,
565 penalty_last_n: i32,
566 penalty_repeat: f32,
567 penalty_freq: f32,
568 penalty_present: f32,
569 ) -> Self {
570 let sampler = unsafe {
571 llama_sampler_init_penalties(
572 n_vocab,
573 penalty_last_n,
574 penalty_repeat,
575 penalty_freq,
576 penalty_present,
577 )
578 };
579 Self {
580 sampler: NonNull::new(sampler).unwrap(),
581 }
582 }
583
584 /// Same as [`Self::penalties`] with sensible defaults:
585 /// `penalty_freq = 0.0` and `penalty_present = 0.0`.
586 ///
587 /// Parameters:
588 /// - `n_vocab`: [`LlamaModel::n_vocab`]
589 /// - `penalty_last_n`: last n tokens to penalize (0 = disable)
590 /// - `penalty_repeat`: repetition penalty (must be > 0.0, 1.0 = disabled)
591 ///
592 /// # Panics
593 ///
594 /// Panics if llama.cpp returns a null pointer.
595 #[must_use]
596 pub fn penalties_simple(n_vocab: i32, penalty_last_n: i32, penalty_repeat: f32) -> Self {
597 Self::penalties(
598 n_vocab,
599 #[allow(clippy::cast_precision_loss)]
600 {
601 penalty_last_n
602 },
603 #[allow(clippy::cast_precision_loss)]
604 {
605 penalty_repeat
606 },
607 #[allow(clippy::cast_precision_loss)]
608 {
609 0.0_f32
610 },
611 #[allow(clippy::cast_precision_loss)]
612 {
613 0.0_f32
614 },
615 )
616 }
617
618 /// Mirostat 1.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
619 ///
620 /// # Panics
621 ///
622 /// Panics if llama.cpp returns a null pointer.
623 ///
624 /// # Parameters:
625 /// - `n_vocab`: [`LlamaModel::n_vocab`]
626 /// - `seed`: Seed to initialize random generation with.
627 /// - `tau`: The target cross-entropy (or surprise) value you want to achieve for the
628 /// generated text. A higher value corresponds to more surprising or less predictable text,
629 /// while a lower value corresponds to less surprising or more predictable text.
630 /// - `eta`: The learning rate used to update `mu` based on the error between the target and
631 /// observed surprisal of the sampled word. A larger learning rate will cause `mu` to be
632 /// updated more quickly, while a smaller learning rate will result in slower updates.
633 /// - `m`: The number of tokens considered in the estimation of `s_hat`. This is an arbitrary
634 /// value that is used to calculate `s_hat`, which in turn helps to calculate the value of `k`.
635 /// In the paper, they use `m = 100`, but you can experiment with different values to see how
636 /// it affects the performance of the algorithm.
637 #[must_use]
638 pub fn mirostat(n_vocab: i32, seed: u32, tau: f32, eta: f32, m: i32) -> Self {
639 let sampler = unsafe { llama_sampler_init_mirostat(n_vocab, seed, tau, eta, m) };
640 Self {
641 sampler: NonNull::new(sampler).unwrap(),
642 }
643 }
644
645 /// Mirostat 2.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
646 ///
647 /// # Panics
648 ///
649 /// Panics if llama.cpp returns a null pointer.
650 ///
651 /// # Parameters:
652 /// - `seed`: Seed to initialize random generation with.
653 /// - `tau`: The target cross-entropy (or surprise) value you want to achieve for the
654 /// generated text. A higher value corresponds to more surprising or less predictable text,
655 /// while a lower value corresponds to less surprising or more predictable text.
656 /// - `eta`: The learning rate used to update `mu` based on the error between the target and
657 /// observed surprisal of the sampled word. A larger learning rate will cause `mu` to be
658 /// updated more quickly, while a smaller learning rate will result in slower updates.
659 #[must_use]
660 pub fn mirostat_v2(seed: u32, tau: f32, eta: f32) -> Self {
661 let sampler = unsafe { llama_sampler_init_mirostat_v2(seed, tau, eta) };
662 Self {
663 sampler: NonNull::new(sampler).unwrap(),
664 }
665 }
666
667 /// Selects a token at random based on each token's probabilities.
668 ///
669 /// # Panics
670 ///
671 /// Panics if llama.cpp returns a null pointer.
672 #[must_use]
673 pub fn dist(seed: u32) -> Self {
674 let sampler = unsafe { llama_sampler_init_dist(seed) };
675 Self {
676 sampler: NonNull::new(sampler).unwrap(),
677 }
678 }
679
680 /// Selects the most likely token.
681 ///
682 /// # Panics
683 ///
684 /// Panics if llama.cpp returns a null pointer.
685 ///
686 /// # Example:
687 /// ```rust
688 /// use llama_cpp_4::token::{
689 /// LlamaToken,
690 /// data::LlamaTokenData,
691 /// data_array::LlamaTokenDataArray
692 /// };
693 /// use llama_cpp_4::sampling::LlamaSampler;
694 ///
695 /// let mut data_array = LlamaTokenDataArray::new(vec![
696 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
697 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
698 /// ], false);
699 ///
700 /// data_array.apply_sampler(&mut LlamaSampler::greedy());
701 ///
702 /// assert_eq!(data_array.data.len(), 2);
703 /// assert_eq!(data_array.selected_token(), Some(LlamaToken(1)));
704 /// ```
705 #[must_use]
706 pub fn greedy() -> Self {
707 let sampler = unsafe { llama_sampler_init_greedy() };
708 Self {
709 sampler: NonNull::new(sampler).unwrap(),
710 }
711 }
712
713 /// Top-N sigma sampling.
714 ///
715 /// Keeps tokens within N standard deviations of the maximum logit.
716 ///
717 /// # Panics
718 ///
719 /// Panics if llama.cpp returns a null pointer.
720 #[must_use]
721 pub fn top_n_sigma(n: f32) -> Self {
722 let sampler = unsafe { llama_sampler_init_top_n_sigma(n) };
723 Self {
724 sampler: NonNull::new(sampler).unwrap(),
725 }
726 }
727
728 /// Adaptive P sampling.
729 ///
730 /// # Panics
731 ///
732 /// Panics if llama.cpp returns a null pointer.
733 ///
734 /// # Parameters
735 /// - `target`: Target probability.
736 /// - `decay`: Decay rate.
737 /// - `seed`: Random seed.
738 #[must_use]
739 pub fn adaptive_p(target: f32, decay: f32, seed: u32) -> Self {
740 let sampler = unsafe { llama_sampler_init_adaptive_p(target, decay, seed) };
741 Self {
742 sampler: NonNull::new(sampler).unwrap(),
743 }
744 }
745
746 /// Logit bias sampler.
747 ///
748 /// Applies additive bias to specific token logits before sampling.
749 ///
750 /// # Panics
751 ///
752 /// Panics if llama.cpp returns a null pointer.
753 ///
754 /// # Parameters
755 /// - `n_vocab`: Number of tokens in the vocabulary ([`LlamaModel::n_vocab`]).
756 /// - `biases`: Slice of `(token_id, bias)` pairs.
757 #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
758 #[must_use]
759 pub fn logit_bias(n_vocab: i32, biases: &[(LlamaToken, f32)]) -> Self {
760 let logit_biases: Vec<llama_logit_bias> = biases
761 .iter()
762 .map(|(token, bias)| llama_logit_bias {
763 token: token.0,
764 bias: *bias,
765 })
766 .collect();
767
768 let sampler = unsafe {
769 llama_sampler_init_logit_bias(n_vocab, logit_biases.len() as i32, logit_biases.as_ptr())
770 };
771 Self {
772 sampler: NonNull::new(sampler).unwrap(),
773 }
774 }
775
776 /// Infill sampler.
777 ///
778 /// Reorders token probabilities for fill-in-the-middle tasks.
779 ///
780 /// # Panics
781 ///
782 /// Panics if llama.cpp returns a null pointer.
783 #[must_use]
784 pub fn infill(model: &LlamaModel) -> Self {
785 let sampler = unsafe { llama_sampler_init_infill(model.get_vocab().vocab.as_ref()) };
786 Self {
787 sampler: NonNull::new(sampler).unwrap(),
788 }
789 }
790
791 /// Get the seed of the sampler.
792 ///
793 /// Returns `LLAMA_DEFAULT_SEED` if the sampler is not seeded.
794 #[must_use]
795 pub fn get_seed(&self) -> u32 {
796 unsafe { llama_sampler_get_seed(self.sampler.as_ptr()) }
797 }
798
799 /// Get the name of the sampler.
800 ///
801 /// # Disabled samplers
802 ///
803 /// When a constructor is handed parameters that make it a no-op — e.g.
804 /// [`Self::temp`] with `1.0`, or [`Self::penalties`] with `penalty_last_n =
805 /// 0` — llama.cpp does not build that sampler. It substitutes an identity
806 /// sampler whose name carries a `?` prefix (`"?temp"`, `"?penalties"`).
807 /// Construction still succeeds, so this name is the only signal that the
808 /// sampler will not do anything. Affects `temp`, `temp_ext`, `top_k`,
809 /// `top_p`, `min_p`, `typical`, `xtc`, `top_n_sigma`, `dry`, and
810 /// `penalties`.
811 ///
812 /// # Panics
813 ///
814 /// Panics if the name is not valid UTF-8.
815 #[must_use]
816 pub fn name(&self) -> String {
817 let c_str = unsafe { llama_sampler_name(self.sampler.as_ptr()) };
818 let c_str = unsafe { std::ffi::CStr::from_ptr(c_str) };
819 c_str
820 .to_str()
821 .expect("sampler name is not valid UTF-8")
822 .to_owned()
823 }
824
825 /// Reset the sampler state (e.g. grammar, repetition penalties).
826 pub fn reset(&mut self) {
827 unsafe { llama_sampler_reset(self.sampler.as_ptr()) }
828 }
829
830 /// Get the number of samplers in a chain.
831 ///
832 /// Returns 0 if this sampler is not a chain.
833 #[must_use]
834 pub fn chain_n(&self) -> i32 {
835 unsafe { llama_sampler_chain_n(self.sampler.as_ptr()) }
836 }
837
838 /// Remove and return the sampler at position `i` from a chain.
839 ///
840 /// The returned sampler is owned by the caller and will be freed on drop.
841 ///
842 /// # Panics
843 ///
844 /// Panics if `i` is out of range or if llama.cpp returns a null pointer.
845 #[must_use]
846 pub fn chain_remove(&mut self, i: i32) -> Self {
847 let sampler = unsafe { llama_sampler_chain_remove(self.sampler.as_ptr(), i) };
848 Self {
849 sampler: NonNull::new(sampler).expect("chain_remove returned null"),
850 }
851 }
852
853 /// Grammar sampler with lazy activation via regex patterns.
854 ///
855 /// The grammar is only activated when one of the trigger patterns or trigger tokens matches.
856 ///
857 /// # Panics
858 /// - If `grammar_str` or `grammar_root` contain null bytes.
859 /// - If any trigger pattern contains null bytes.
860 /// - If llama.cpp returns a null pointer.
861 ///
862 /// Prefer [`Self::try_grammar_lazy_patterns`] when any of those strings came
863 /// from a request rather than your own source.
864 #[must_use]
865 pub fn grammar_lazy_patterns(
866 model: &LlamaModel,
867 grammar_str: &str,
868 grammar_root: &str,
869 trigger_patterns: &[&str],
870 trigger_tokens: &[LlamaToken],
871 ) -> Self {
872 Self::try_grammar_lazy_patterns(
873 model,
874 grammar_str,
875 grammar_root,
876 trigger_patterns,
877 trigger_tokens,
878 )
879 .unwrap_or_else(|e| panic!("building lazy grammar sampler: {e}"))
880 }
881
882 /// [`Self::grammar_lazy_patterns`], reporting failures instead of panicking.
883 ///
884 /// # Errors
885 ///
886 /// [`SamplerInitError::InteriorNul`] for a NUL in the grammar, the rule name
887 /// or any trigger pattern, or [`SamplerInitError::NullReturn`] if llama.cpp
888 /// rejected the grammar.
889 pub fn try_grammar_lazy_patterns(
890 model: &LlamaModel,
891 grammar_str: &str,
892 grammar_root: &str,
893 trigger_patterns: &[&str],
894 trigger_tokens: &[LlamaToken],
895 ) -> Result<Self, SamplerInitError> {
896 let grammar_str = c_string("grammar_str", grammar_str)?;
897 let grammar_root = c_string("grammar_root", grammar_root)?;
898 let pattern_cstrings = trigger_patterns
899 .iter()
900 .map(|pattern| c_string("trigger_patterns", pattern))
901 .collect::<Result<Vec<CString>, _>>()?;
902 let mut pattern_ptrs: Vec<*const c_char> =
903 pattern_cstrings.iter().map(|s| s.as_ptr()).collect();
904
905 // SAFETY: every pointer is borrowed from a `CString`/slice that outlives
906 // the call, and the two lengths describe exactly those buffers.
907 // `trigger_tokens` is `#[repr(transparent)]` over `llama_token`, so the
908 // cast is a layout no-op. llama.cpp copies what it needs.
909 let sampler = unsafe {
910 llama_sampler_init_grammar_lazy_patterns(
911 model.get_vocab().vocab.as_ref(),
912 grammar_str.as_ptr(),
913 grammar_root.as_ptr(),
914 pattern_ptrs.as_mut_ptr(),
915 pattern_ptrs.len(),
916 trigger_tokens.as_ptr().cast(),
917 trigger_tokens.len(),
918 )
919 };
920 NonNull::new(sampler)
921 .map(|sampler| Self { sampler })
922 .ok_or(SamplerInitError::NullReturn("lazy grammar"))
923 }
924
925 /// Clone this sampler.
926 ///
927 /// Creates an independent copy of this sampler with the same state.
928 ///
929 /// # Panics
930 ///
931 /// Panics if llama.cpp returns a null pointer.
932 #[must_use]
933 pub fn clone_sampler(&self) -> Self {
934 let sampler = unsafe { llama_sampler_clone(self.sampler.as_ptr()) };
935 Self {
936 sampler: NonNull::new(sampler).expect("sampler_clone returned null"),
937 }
938 }
939
940 /// Copy mutable state from `src` into this sampler, in place.
941 ///
942 /// Unlike [`Self::clone_sampler`], which allocates a new sampler, this
943 /// overwrites the state of an existing one and so is the cheap way to
944 /// rewind a sampler to a checkpoint in a loop. Added upstream in
945 /// llama.cpp `b10470` (`llama_sampler_copy`).
946 ///
947 /// # Safety and preconditions
948 ///
949 /// llama.cpp requires `src` and `self` to be **the same sampler type with
950 /// the same configuration** — e.g. two `dist` samplers, or two chains built
951 /// the same way. Copying between mismatched samplers is undefined
952 /// behaviour upstream and is not checked here, so treat the pairing as the
953 /// caller's contract. A sampler produced by `src.clone_sampler()` always
954 /// satisfies it.
955 pub fn copy_state_from(&mut self, src: &Self) {
956 unsafe { llama_sampler_copy(src.sampler.as_ptr(), self.sampler.as_ptr()) }
957 }
958
959 /// Print sampler performance data.
960 pub fn perf_print(&self) {
961 unsafe { llama_cpp_sys_4::llama_perf_sampler_print(self.sampler.as_ptr()) }
962 }
963
964 /// Reset sampler performance counters.
965 pub fn perf_reset(&mut self) {
966 unsafe { llama_cpp_sys_4::llama_perf_sampler_reset(self.sampler.as_ptr()) }
967 }
968
969 /// Get sampler performance data.
970 #[must_use]
971 pub fn perf_data(&self) -> llama_cpp_sys_4::llama_perf_sampler_data {
972 unsafe { llama_cpp_sys_4::llama_perf_sampler(self.sampler.as_ptr()) }
973 }
974
975 /// Get a non-owning reference to the `i`th sampler in a chain.
976 ///
977 /// # Safety
978 ///
979 /// The returned pointer is owned by the chain. Do not free it or use it
980 /// after the chain is dropped or modified.
981 #[must_use]
982 pub unsafe fn chain_get_ptr(&self, i: i32) -> *mut llama_sampler {
983 llama_cpp_sys_4::llama_sampler_chain_get(self.sampler.as_ptr(), i)
984 }
985
986 /// Create a sampler from a raw interface and context.
987 ///
988 /// # Safety
989 ///
990 /// The caller must ensure that `iface` and `ctx` are valid and that the
991 /// interface functions properly manage the context lifetime.
992 ///
993 /// # Panics
994 ///
995 /// Panics if llama.cpp returns a null pointer.
996 #[must_use]
997 pub unsafe fn from_raw(
998 iface: *mut llama_cpp_sys_4::llama_sampler_i,
999 ctx: llama_cpp_sys_4::llama_sampler_context_t,
1000 ) -> Self {
1001 let sampler = llama_cpp_sys_4::llama_sampler_init(iface, ctx);
1002 Self {
1003 sampler: NonNull::new(sampler).expect("sampler_init returned null"),
1004 }
1005 }
1006
1007 /// Adopt a raw `llama_sampler *`, taking ownership.
1008 ///
1009 /// # Safety
1010 ///
1011 /// `raw` must come from a llama.cpp entry point documented as returning a
1012 /// sampler the caller owns and releases with `llama_sampler_free`, and must
1013 /// not be owned by anything else — [`Drop`] frees it.
1014 pub(crate) unsafe fn from_raw_ptr(raw: NonNull<llama_sampler>) -> Self {
1015 Self { sampler: raw }
1016 }
1017
1018 /// The underlying `llama_sampler *`. Borrowed; still owned by `self`.
1019 pub(crate) fn as_ptr(&self) -> *mut llama_sampler {
1020 self.sampler.as_ptr()
1021 }
1022
1023 /// Creates a new instance of `LlamaSampler` with common sampling parameters.
1024 ///
1025 /// This function initializes a `LlamaSampler` using default values from `common_sampler_params`
1026 /// and configures it with common settings such as `top_k`, `top_p`, `temperature`, and `seed` values.
1027 ///
1028 /// # Panics
1029 ///
1030 /// Panics if llama.cpp returns a null pointer.
1031 ///
1032 /// # Returns
1033 /// A `LlamaSampler` instance configured with the common sampling parameters.
1034 #[must_use]
1035 pub fn common() -> Self {
1036 let params = common_sampler_params::default();
1037
1038 let sampler = unsafe {
1039 let mut sparams = llama_sampler_chain_default_params();
1040 sparams.no_perf = false;
1041
1042 let smpl = llama_sampler_chain_init(sparams);
1043
1044 llama_sampler_chain_add(smpl, llama_sampler_init_top_k(params.top_k));
1045 llama_sampler_chain_add(
1046 smpl,
1047 #[allow(clippy::cast_sign_loss)]
1048 llama_sampler_init_top_p(params.top_p, params.min_keep as usize),
1049 );
1050 llama_sampler_chain_add(smpl, llama_sampler_init_temp(params.temp));
1051 #[allow(clippy::cast_sign_loss)]
1052 llama_sampler_chain_add(smpl, llama_sampler_init_dist(params.seed));
1053
1054 smpl
1055 };
1056
1057 Self {
1058 sampler: NonNull::new(sampler).unwrap(),
1059 }
1060 }
1061}
1062
1063impl Drop for LlamaSampler {
1064 fn drop(&mut self) {
1065 unsafe {
1066 llama_sampler_free(self.sampler.as_ptr());
1067 }
1068 }
1069}