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