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