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