llama_cpp_2/sampling.rs
1//! Safe wrapper around `llama_sampler`.
2
3use std::borrow::Borrow;
4use std::ffi::{c_char, c_int, CStr, CString};
5use std::fmt;
6use std::mem::ManuallyDrop;
7
8use crate::context::LlamaContext;
9use crate::model::LlamaModel;
10use crate::ptr::Ptr;
11#[cfg(feature = "common")]
12use crate::status_is_ok;
13use crate::token::data_array::LlamaTokenDataArray;
14use crate::token::logit_bias::LlamaLogitBias;
15use crate::token::LlamaToken;
16use crate::{GrammarError, SamplerAcceptError};
17
18/// A safe wrapper around `llama_sampler`.
19pub struct LlamaSampler {
20 pub(crate) sampler: Ptr<llama_cpp_sys_2::llama_sampler>,
21}
22
23// SAFETY: Samplers are safe to clone and drop from any thread.
24//
25// (It is not safe to _use_ them from multiple threads at a time, but we
26// disallow that by using `&mut` in those methods).
27unsafe impl Send for LlamaSampler {}
28// SAFETY: Same as above.
29unsafe impl Sync for LlamaSampler {}
30
31impl fmt::Debug for LlamaSampler {
32 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33 let mut f = f.debug_struct("LlamaSampler");
34
35 f.field("name", &self.name().unwrap_or(c"(NULL)"));
36
37 if let Some(chain) = self.as_chain() {
38 f.field("chain", &chain);
39 }
40
41 f.finish()
42 }
43}
44
45impl Clone for LlamaSampler {
46 fn clone(&self) -> Self {
47 let ptr = unsafe { llama_cpp_sys_2::llama_sampler_clone(self.sampler.as_ptr()) };
48 let sampler = Ptr::new(ptr).expect("failed cloning sampler");
49 Self { sampler }
50 }
51}
52
53impl LlamaSampler {
54 #[inline]
55 pub(crate) fn from_raw(ptr: *mut llama_cpp_sys_2::llama_sampler) -> Option<Self> {
56 let sampler = Ptr::new(ptr)?;
57 Some(Self { sampler })
58 }
59
60 /// Returns `Some(...)` if the sampler is a `llama_sampler_chain`.
61 #[inline]
62 fn as_chain(&self) -> Option<LlamaSamplerChain<'_>> {
63 let chain = unsafe {
64 llama_cpp_sys_2::llama_sampler_chain_get(self.sampler.as_mut_ptr_unsound(), -1)
65 };
66
67 if chain.is_null() {
68 None
69 } else {
70 Some(LlamaSamplerChain(self))
71 }
72 }
73
74 /// The name of the sampler, if any.
75 #[inline]
76 pub fn name(&self) -> Option<&CStr> {
77 let ptr = unsafe { llama_cpp_sys_2::llama_sampler_name(self.sampler.as_ptr()) };
78 if ptr.is_null() {
79 None
80 } else {
81 Some(unsafe { CStr::from_ptr(ptr) })
82 }
83 }
84
85 /// Sample and accept a token from the idx-th output of the last evaluation
86 #[must_use]
87 pub fn sample(&mut self, ctx: &LlamaContext<'_>, idx: i32) -> LlamaToken {
88 let token = unsafe {
89 llama_cpp_sys_2::llama_sampler_sample(
90 self.sampler.as_mut_ptr(),
91 ctx.synchronizable_ptr(),
92 idx,
93 )
94 };
95
96 LlamaToken(token)
97 }
98
99 /// Applies this sampler to a [`LlamaTokenDataArray`].
100 pub fn apply(&mut self, data_array: &mut LlamaTokenDataArray) {
101 data_array.apply_sampler(self);
102 }
103
104 /// Accepts a token from the sampler, possibly updating the internal state of certain samplers
105 /// (e.g. grammar, repetition, etc.)
106 pub fn accept(&mut self, token: LlamaToken) {
107 #[cfg(feature = "common")]
108 {
109 let _ = self.try_accept(token);
110 }
111 #[cfg(not(feature = "common"))]
112 unsafe {
113 llama_cpp_sys_2::llama_sampler_accept(self.sampler.as_mut_ptr(), token.0);
114 }
115 }
116
117 /// Accepts several tokens from the sampler or context, possibly updating the internal state of
118 /// certain samplers (e.g. grammar, repetition, etc.)
119 pub fn accept_many(&mut self, tokens: impl IntoIterator<Item = impl Borrow<LlamaToken>>) {
120 for token in tokens {
121 self.accept(*token.borrow());
122 }
123 }
124
125 /// Accepts several tokens from the sampler or context, possibly updating the internal state of
126 /// certain samplers (e.g. grammar, repetition, etc.)
127 #[must_use]
128 pub fn with_tokens(
129 mut self,
130 tokens: impl IntoIterator<Item = impl Borrow<LlamaToken>>,
131 ) -> Self {
132 self.accept_many(tokens);
133 self
134 }
135
136 /// Try accepting a token from the sampler. Returns an error if the sampler throws.
137 #[cfg(feature = "common")]
138 pub fn try_accept(&mut self, token: LlamaToken) -> Result<(), SamplerAcceptError> {
139 let sampler_result =
140 unsafe { llama_cpp_sys_2::llama_rs_sampler_accept(self.sampler.as_mut_ptr(), token.0) };
141 if status_is_ok(sampler_result) {
142 Ok(())
143 } else {
144 Err(SamplerAcceptError::FfiError(sampler_result))
145 }
146 }
147
148 /// Resets the internal state of the sampler.
149 ///
150 /// This can be useful when you want to start fresh with a sampler without creating a new instance.
151 pub fn reset(&mut self) {
152 unsafe {
153 llama_cpp_sys_2::llama_sampler_reset(self.sampler.as_mut_ptr());
154 }
155 }
156
157 /// Gets the random seed used by this sampler.
158 ///
159 /// Returns:
160 /// - For random samplers (dist, mirostat, mirostat_v2): returns their current seed
161 /// - For sampler chains: returns the first non-default seed found in reverse order
162 /// - For all other samplers: returns 0xFFFFFFFF
163 #[must_use]
164 pub fn get_seed(&self) -> u32 {
165 unsafe { llama_cpp_sys_2::llama_sampler_get_seed(self.sampler.as_ptr()) }
166 }
167
168 /// Combines a list of samplers into a single sampler that applies each component sampler one
169 /// after another.
170 ///
171 /// If you are using a chain to select a token, the chain should always end with one of
172 /// [`LlamaSampler::greedy`], [`LlamaSampler::dist`], [`LlamaSampler::mirostat`], and
173 /// [`LlamaSampler::mirostat_v2`].
174 #[must_use]
175 pub fn chain(samplers: impl IntoIterator<Item = Self>, no_perf: bool) -> Self {
176 let params = llama_cpp_sys_2::llama_sampler_chain_params { no_perf };
177
178 let chain = unsafe { llama_cpp_sys_2::llama_sampler_chain_init(params) };
179 let mut chain = Ptr::new(chain).expect("failed creating sampler chain");
180
181 for sampler in samplers {
182 // Do not call `llama_sampler_free` on the sampler, we want to
183 // transfer ownership to the chain.
184 let mut sampler = ManuallyDrop::new(sampler);
185
186 unsafe {
187 llama_cpp_sys_2::llama_sampler_chain_add(
188 chain.as_mut_ptr(),
189 sampler.sampler.as_mut_ptr(),
190 )
191 };
192 }
193
194 Self { sampler: chain }
195 }
196
197 /// Same as [`Self::chain`] with `no_perf = false`.
198 ///
199 /// # Example
200 /// ```rust
201 /// use llama_cpp_2::token::{
202 /// LlamaToken,
203 /// data::LlamaTokenData,
204 /// data_array::LlamaTokenDataArray
205 /// };
206 /// use llama_cpp_2::sampling::LlamaSampler;
207 /// use llama_cpp_2::llama_backend::LlamaBackend;
208 /// let backend = LlamaBackend::init().unwrap();
209 ///
210 /// let mut data_array = LlamaTokenDataArray::new(vec![
211 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
212 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
213 /// LlamaTokenData::new(LlamaToken(2), 2., 0.),
214 /// ], false);
215 ///
216 /// data_array.apply_sampler(&mut LlamaSampler::chain_simple([
217 /// LlamaSampler::temp(0.5),
218 /// LlamaSampler::greedy(),
219 /// ]));
220 ///
221 /// assert_eq!(data_array.data[0].logit(), 0.);
222 /// assert_eq!(data_array.data[1].logit(), 2.);
223 /// assert_eq!(data_array.data[2].logit(), 4.);
224 ///
225 /// assert_eq!(data_array.data.len(), 3);
226 /// assert_eq!(data_array.selected_token(), Some(LlamaToken(2)));
227 /// ```
228 #[must_use]
229 pub fn chain_simple(samplers: impl IntoIterator<Item = Self>) -> Self {
230 Self::chain(samplers, false)
231 }
232
233 #[allow(clippy::doc_markdown)]
234 /// Updates the logits l_i' = l_i/t. When t <= 0.0f, the maximum logit is kept at it's original
235 /// value, the rest are set to -inf
236 ///
237 /// # Example:
238 /// ```rust
239 /// use llama_cpp_2::token::{
240 /// LlamaToken,
241 /// data::LlamaTokenData,
242 /// data_array::LlamaTokenDataArray
243 /// };
244 /// use llama_cpp_2::sampling::LlamaSampler;
245 ///
246 /// let mut data_array = LlamaTokenDataArray::new(vec![
247 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
248 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
249 /// LlamaTokenData::new(LlamaToken(2), 2., 0.),
250 /// ], false);
251 ///
252 /// data_array.apply_sampler(&mut LlamaSampler::temp(0.5));
253 ///
254 /// assert_eq!(data_array.data[0].logit(), 0.);
255 /// assert_eq!(data_array.data[1].logit(), 2.);
256 /// assert_eq!(data_array.data[2].logit(), 4.);
257 /// ```
258 #[must_use]
259 pub fn temp(t: f32) -> Self {
260 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_temp(t) };
261 Self::from_raw(sampler).expect("failed creating sampler")
262 }
263
264 /// Dynamic temperature implementation (a.k.a. entropy) described in the paper
265 /// <https://arxiv.org/abs/2309.02772>.
266 #[must_use]
267 pub fn temp_ext(t: f32, delta: f32, exponent: f32) -> Self {
268 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_temp_ext(t, delta, exponent) };
269 Self::from_raw(sampler).expect("failed creating sampler")
270 }
271
272 /// Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration"
273 /// <https://arxiv.org/abs/1904.09751>
274 ///
275 /// # Example:
276 /// ```rust
277 /// use llama_cpp_2::token::{
278 /// LlamaToken,
279 /// data::LlamaTokenData,
280 /// data_array::LlamaTokenDataArray
281 /// };
282 /// use llama_cpp_2::sampling::LlamaSampler;
283 ///
284 /// let mut data_array = LlamaTokenDataArray::new(vec![
285 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
286 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
287 /// LlamaTokenData::new(LlamaToken(2), 2., 0.),
288 /// LlamaTokenData::new(LlamaToken(3), 3., 0.),
289 /// ], false);
290 ///
291 /// data_array.apply_sampler(&mut LlamaSampler::top_k(2));
292 ///
293 /// assert_eq!(data_array.data.len(), 2);
294 /// assert_eq!(data_array.data[0].id(), LlamaToken(3));
295 /// assert_eq!(data_array.data[1].id(), LlamaToken(2));
296 /// ```
297 #[must_use]
298 pub fn top_k(k: i32) -> Self {
299 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_top_k(k) };
300 Self::from_raw(sampler).expect("failed creating sampler")
301 }
302
303 /// Top-nσ sampling as described in academic paper "Top-nσ: Not All Logits Are You Need"
304 /// <https://arxiv.org/pdf/2411.07641>
305 ///
306 /// This method filters logits by selecting only those within *n* standard deviations of the mean.
307 ///
308 /// # Parameters
309 /// - `n`: Number of standard deviations from the mean to include in sampling
310 ///
311 /// # Example
312 /// ```rust
313 /// use llama_cpp_2::sampling::LlamaSampler;
314 /// use llama_cpp_2::token::{
315 /// LlamaToken,
316 /// data::LlamaTokenData,
317 /// data_array::LlamaTokenDataArray
318 /// };
319 ///
320 /// let mut data_array = LlamaTokenDataArray::new(vec![
321 /// LlamaTokenData::new(LlamaToken(0), 0.0, 0.0),
322 /// LlamaTokenData::new(LlamaToken(1), 1.0, 0.0),
323 /// LlamaTokenData::new(LlamaToken(2), 2.0, 0.0),
324 /// ], false);
325 ///
326 /// data_array.apply_sampler(&mut LlamaSampler::top_n_sigma(2.0));
327 /// ```
328 #[must_use]
329 pub fn top_n_sigma(n: f32) -> Self {
330 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_top_n_sigma(n) };
331 Self::from_raw(sampler).expect("failed creating sampler")
332 }
333
334 /// Locally Typical Sampling implementation described in the paper <https://arxiv.org/abs/2202.00666>.
335 #[must_use]
336 pub fn typical(p: f32, min_keep: usize) -> Self {
337 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_typical(p, min_keep) };
338 Self::from_raw(sampler).expect("failed creating sampler")
339 }
340
341 /// Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration"
342 /// <https://arxiv.org/abs/1904.09751>
343 #[must_use]
344 pub fn top_p(p: f32, min_keep: usize) -> Self {
345 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_top_p(p, min_keep) };
346 Self::from_raw(sampler).expect("failed creating sampler")
347 }
348
349 /// Minimum P sampling as described in <https://github.com/ggerganov/llama.cpp/pull/3841>
350 #[must_use]
351 pub fn min_p(p: f32, min_keep: usize) -> Self {
352 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_min_p(p, min_keep) };
353 Self::from_raw(sampler).expect("failed creating sampler")
354 }
355
356 /// XTC sampler as described in <https://github.com/oobabooga/text-generation-webui/pull/6335>
357 #[must_use]
358 pub fn xtc(p: f32, t: f32, min_keep: usize, seed: u32) -> Self {
359 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_xtc(p, t, min_keep, seed) };
360 Self::from_raw(sampler).expect("failed creating sampler")
361 }
362
363 /// Grammar sampler
364 #[must_use]
365 pub fn grammar(
366 model: &LlamaModel,
367 grammar_str: &str,
368 grammar_root: &str,
369 ) -> Result<Self, GrammarError> {
370 let (grammar_str, grammar_root) =
371 Self::sanitize_grammar_strings(grammar_str, grammar_root)?;
372
373 #[cfg(feature = "common")]
374 let sampler = unsafe {
375 llama_cpp_sys_2::llama_rs_sampler_init_grammar(
376 model.vocab().as_ptr(),
377 grammar_str.as_ptr(),
378 grammar_root.as_ptr(),
379 )
380 };
381 #[cfg(not(feature = "common"))]
382 let sampler = unsafe {
383 llama_cpp_sys_2::llama_sampler_init_grammar(
384 model.vocab().as_ptr(),
385 grammar_str.as_ptr(),
386 grammar_root.as_ptr(),
387 )
388 };
389
390 if sampler.is_null() {
391 Err(GrammarError::NullGrammar)
392 } else {
393 Ok(Self::from_raw(sampler).unwrap())
394 }
395 }
396
397 /// Lazy grammar sampler, introduced in <https://github.com/ggerganov/llama.cpp/pull/9639>
398 ///
399 /// This sampler enforces grammar rules only when specific trigger words or tokens are encountered.
400 #[cfg(feature = "common")]
401 #[must_use]
402 pub fn grammar_lazy(
403 model: &LlamaModel,
404 grammar_str: &str,
405 grammar_root: &str,
406 trigger_words: impl IntoIterator<Item = impl AsRef<[u8]>>,
407 trigger_tokens: &[LlamaToken],
408 ) -> Result<Self, GrammarError> {
409 let (grammar_str, grammar_root) =
410 Self::sanitize_grammar_strings(grammar_str, grammar_root)?;
411 let trigger_words = Self::sanitize_trigger_words(trigger_words)?;
412
413 let mut trigger_word_ptrs: Vec<*const c_char> =
414 trigger_words.iter().map(|cs| cs.as_ptr()).collect();
415
416 let sampler = unsafe {
417 llama_cpp_sys_2::llama_rs_sampler_init_grammar_lazy(
418 model.vocab().as_ptr(),
419 grammar_str.as_ptr(),
420 grammar_root.as_ptr(),
421 trigger_word_ptrs.as_mut_ptr(),
422 trigger_word_ptrs.len(),
423 trigger_tokens.as_ptr().cast(),
424 trigger_tokens.len(),
425 )
426 };
427
428 if sampler.is_null() {
429 Err(GrammarError::NullGrammar)
430 } else {
431 Ok(Self::from_raw(sampler).unwrap())
432 }
433 }
434
435 /// Lazy grammar sampler using regex trigger patterns.
436 ///
437 /// Trigger patterns are regular expressions matched from the start of the
438 /// generation output. The grammar sampler will be fed content starting from
439 /// the first match group.
440 ///
441 /// Without the `common` feature there is no C++ shim to catch exceptions,
442 /// and `llama.cpp` builds a `std::regex` from each trigger pattern, so an
443 /// invalid pattern aborts instead of returning
444 /// [`GrammarError::NullGrammar`]. An unparseable grammar is reported with a
445 /// null pointer either way.
446 #[must_use]
447 pub fn grammar_lazy_patterns(
448 model: &LlamaModel,
449 grammar_str: &str,
450 grammar_root: &str,
451 trigger_patterns: &[String],
452 trigger_tokens: &[LlamaToken],
453 ) -> Result<Self, GrammarError> {
454 let (grammar_str, grammar_root) =
455 Self::sanitize_grammar_strings(grammar_str, grammar_root)?;
456 let trigger_patterns = Self::sanitize_trigger_patterns(trigger_patterns)?;
457
458 let mut trigger_pattern_ptrs: Vec<*const c_char> =
459 trigger_patterns.iter().map(|cs| cs.as_ptr()).collect();
460
461 #[cfg(feature = "common")]
462 let sampler = unsafe {
463 llama_cpp_sys_2::llama_rs_sampler_init_grammar_lazy_patterns(
464 model.vocab().as_ptr(),
465 grammar_str.as_ptr(),
466 grammar_root.as_ptr(),
467 trigger_pattern_ptrs.as_mut_ptr(),
468 trigger_pattern_ptrs.len(),
469 trigger_tokens.as_ptr().cast(),
470 trigger_tokens.len(),
471 )
472 };
473 #[cfg(not(feature = "common"))]
474 let sampler = unsafe {
475 llama_cpp_sys_2::llama_sampler_init_grammar_lazy_patterns(
476 model.vocab().as_ptr(),
477 grammar_str.as_ptr(),
478 grammar_root.as_ptr(),
479 trigger_pattern_ptrs.as_mut_ptr(),
480 trigger_pattern_ptrs.len(),
481 trigger_tokens.as_ptr().cast(),
482 trigger_tokens.len(),
483 )
484 };
485
486 if sampler.is_null() {
487 Err(GrammarError::NullGrammar)
488 } else {
489 Ok(Self::from_raw(sampler).unwrap())
490 }
491 }
492
493 /// Builds the `toktrie` tokenizer environment for `model`.
494 ///
495 /// Use this to construct your own `llguidance::ParserFactory` (with any slice
496 /// regexes, inference capabilities, or other configuration llguidance supports) and,
497 /// from it, a `llguidance::Matcher` to convert into a `LlamaSampler` via
498 /// `LlamaSampler::from` (or `.into()`).
499 #[cfg(feature = "llguidance")]
500 #[must_use]
501 pub fn llguidance_tok_env(model: &LlamaModel) -> toktrie::TokEnv {
502 crate::llguidance_sampler::llguidance_build_tok_env(model)
503 }
504
505 fn sanitize_grammar_strings(
506 grammar_str: &str,
507 grammar_root: &str,
508 ) -> Result<(CString, CString), GrammarError> {
509 if !grammar_str.contains(grammar_root) {
510 return Err(GrammarError::RootNotFound);
511 }
512
513 if grammar_str.contains('\0') || grammar_root.contains('\0') {
514 return Err(GrammarError::GrammarNullBytes);
515 }
516
517 Ok((
518 CString::new(grammar_str).unwrap(),
519 CString::new(grammar_root).unwrap(),
520 ))
521 }
522
523 fn sanitize_trigger_words(
524 trigger_words: impl IntoIterator<Item = impl AsRef<[u8]>>,
525 ) -> Result<Vec<CString>, GrammarError> {
526 let trigger_words: Vec<_> = trigger_words.into_iter().collect();
527 if trigger_words
528 .iter()
529 .any(|word| word.as_ref().contains(&b'\0'))
530 {
531 return Err(GrammarError::TriggerWordNullBytes);
532 }
533 Ok(trigger_words
534 .into_iter()
535 .map(|word| CString::new(word.as_ref()).unwrap())
536 .collect())
537 }
538
539 fn sanitize_trigger_patterns(
540 trigger_patterns: &[String],
541 ) -> Result<Vec<CString>, GrammarError> {
542 let mut patterns = Vec::with_capacity(trigger_patterns.len());
543 for pattern in trigger_patterns {
544 if pattern.contains('\0') {
545 return Err(GrammarError::GrammarNullBytes);
546 }
547 patterns.push(CString::new(pattern.as_str()).unwrap());
548 }
549 Ok(patterns)
550 }
551
552 /// DRY sampler, designed by p-e-w, as described in:
553 /// <https://github.com/oobabooga/text-generation-webui/pull/5677>, porting Koboldcpp
554 /// implementation authored by pi6am: <https://github.com/LostRuins/koboldcpp/pull/982>
555 ///
556 /// # Panics
557 /// If any string in ``seq_breakers`` contains null bytes.
558 #[allow(missing_docs)]
559 #[must_use]
560 pub fn dry(
561 model: &LlamaModel,
562 multiplier: f32,
563 base: f32,
564 allowed_length: i32,
565 penalty_last_n: i32,
566 seq_breakers: impl IntoIterator<Item = impl AsRef<[u8]>>,
567 ) -> Self {
568 let seq_breakers: Vec<CString> = seq_breakers
569 .into_iter()
570 .map(|s| CString::new(s.as_ref()).expect("A sequence breaker contains null bytes"))
571 .collect();
572 let mut seq_breaker_pointers: Vec<*const c_char> =
573 seq_breakers.iter().map(|s| s.as_ptr()).collect();
574
575 let sampler = unsafe {
576 llama_cpp_sys_2::llama_sampler_init_dry(
577 model.vocab().as_ptr(),
578 multiplier,
579 base,
580 allowed_length,
581 penalty_last_n,
582 seq_breaker_pointers.as_mut_ptr(),
583 seq_breaker_pointers.len(),
584 )
585 };
586 Self::from_raw(sampler).expect("failed creating sampler")
587 }
588
589 /// Penalizes tokens for being present in the context.
590 ///
591 /// Parameters:
592 /// - ``n_vocab``: [`LlamaModel::n_vocab`]
593 /// - ``penalty_last_n``: last n tokens to penalize (0 = disable penalty). Negative values are
594 /// clamped to 0; they no longer select the context size.
595 /// - ``penalty_repeat``: must be > 0.0, 1.0 = disabled
596 /// - ``penalty_freq``: must be finite, 0.0 = disabled
597 /// - ``penalty_present``: must be finite, 0.0 = disabled
598 #[allow(clippy::too_many_arguments)]
599 #[must_use]
600 pub fn penalties(
601 n_vocab: i32,
602 penalty_last_n: i32,
603 penalty_repeat: f32,
604 penalty_freq: f32,
605 penalty_present: f32,
606 ) -> Self {
607 let sampler = unsafe {
608 llama_cpp_sys_2::llama_sampler_init_penalties(
609 n_vocab,
610 penalty_last_n,
611 penalty_repeat,
612 penalty_freq,
613 penalty_present,
614 )
615 };
616 Self::from_raw(sampler).expect("failed creating sampler")
617 }
618
619 /// Mirostat 1.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
620 ///
621 /// # Parameters:
622 /// - ``n_vocab``: [`LlamaModel::n_vocab`]
623 /// - ``seed``: Seed to initialize random generation with.
624 /// - ``tau``: The target cross-entropy (or surprise) value you want to achieve for the
625 /// generated text. A higher value corresponds to more surprising or less predictable text,
626 /// while a lower value corresponds to less surprising or more predictable text.
627 /// - ``eta``: The learning rate used to update `mu` based on the error between the target and
628 /// observed surprisal of the sampled word. A larger learning rate will cause `mu` to be
629 /// updated more quickly, while a smaller learning rate will result in slower updates.
630 /// - ``m``: The number of tokens considered in the estimation of `s_hat`. This is an arbitrary
631 /// value that is used to calculate `s_hat`, which in turn helps to calculate the value of `k`.
632 /// In the paper, they use `m = 100`, but you can experiment with different values to see how
633 /// it affects the performance of the algorithm.
634 #[must_use]
635 pub fn mirostat(n_vocab: i32, seed: u32, tau: f32, eta: f32, m: i32) -> Self {
636 let sampler =
637 unsafe { llama_cpp_sys_2::llama_sampler_init_mirostat(n_vocab, seed, tau, eta, m) };
638 Self::from_raw(sampler).expect("failed creating sampler")
639 }
640
641 /// Mirostat 2.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
642 ///
643 /// # Parameters:
644 /// - ``seed``: Seed to initialize random generation with.
645 /// - ``tau``: The target cross-entropy (or surprise) value you want to achieve for the
646 /// generated text. A higher value corresponds to more surprising or less predictable text,
647 /// while a lower value corresponds to less surprising or more predictable text.
648 /// - ``eta``: The learning rate used to update `mu` based on the error between the target and
649 /// observed surprisal of the sampled word. A larger learning rate will cause `mu` to be
650 /// updated more quickly, while a smaller learning rate will result in slower updates.
651 #[must_use]
652 pub fn mirostat_v2(seed: u32, tau: f32, eta: f32) -> Self {
653 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_mirostat_v2(seed, tau, eta) };
654 Self::from_raw(sampler).expect("failed creating sampler")
655 }
656
657 /// Selects a token at random based on each token's probabilities
658 #[must_use]
659 pub fn dist(seed: u32) -> Self {
660 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_dist(seed) };
661 Self::from_raw(sampler).expect("failed creating sampler")
662 }
663
664 /// Selects the most likely token
665 ///
666 /// # Example:
667 /// ```rust
668 /// use llama_cpp_2::token::{
669 /// LlamaToken,
670 /// data::LlamaTokenData,
671 /// data_array::LlamaTokenDataArray
672 /// };
673 /// use llama_cpp_2::sampling::LlamaSampler;
674 ///
675 /// let mut data_array = LlamaTokenDataArray::new(vec![
676 /// LlamaTokenData::new(LlamaToken(0), 0., 0.),
677 /// LlamaTokenData::new(LlamaToken(1), 1., 0.),
678 /// ], false);
679 ///
680 /// data_array.apply_sampler(&mut LlamaSampler::greedy());
681 ///
682 /// assert_eq!(data_array.data.len(), 2);
683 /// assert_eq!(data_array.selected_token(), Some(LlamaToken(1)));
684 /// ```
685 #[must_use]
686 pub fn greedy() -> Self {
687 let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_greedy() };
688 Self::from_raw(sampler).expect("failed creating sampler")
689 }
690
691 /// Creates a sampler that applies bias values to specific tokens during sampling.
692 ///
693 /// # Parameters
694 /// - ``n_vocab``: [`LlamaModel::n_vocab`]
695 /// - ``biases``: Slice of [`LlamaLogitBias`] values specifying token-bias pairs
696 ///
697 /// # Example
698 /// ```rust
699 /// use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
700 /// use llama_cpp_2::sampling::LlamaSampler;
701 ///
702 /// let biases = vec![
703 /// LlamaLogitBias::new(LlamaToken(1), 1.5), // Increase probability of token 1
704 /// LlamaLogitBias::new(LlamaToken(2), -1.0), // Decrease probability of token 2
705 /// ];
706 ///
707 /// // Assuming vocab_size of 32000
708 /// let sampler = LlamaSampler::logit_bias(32000, &biases);
709 /// ```
710 #[must_use]
711 pub fn logit_bias(n_vocab: i32, biases: &[LlamaLogitBias]) -> Self {
712 let data = biases.as_ptr().cast::<llama_cpp_sys_2::llama_logit_bias>();
713
714 let sampler = unsafe {
715 llama_cpp_sys_2::llama_sampler_init_logit_bias(n_vocab, biases.len() as i32, data)
716 };
717
718 Self::from_raw(sampler).expect("failed creating sampler")
719 }
720}
721
722impl Drop for LlamaSampler {
723 fn drop(&mut self) {
724 unsafe { llama_cpp_sys_2::llama_sampler_free(self.sampler.as_mut_ptr()) };
725 }
726}
727
728/// Small helper for iterating over `llama_sampler_chain`.
729struct LlamaSamplerChain<'a>(&'a LlamaSampler);
730
731impl fmt::Debug for LlamaSamplerChain<'_> {
732 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
733 let mut f = f.debug_list();
734 for i in 0..self.n() {
735 // SAFETY: The sampler is only used in this scope.
736 let sampler = unsafe { self.get(i) }.expect("sampler chain must have sampler");
737 f.entry(&&*sampler);
738 }
739 f.finish()
740 }
741}
742
743impl LlamaSamplerChain<'_> {
744 fn n(&self) -> c_int {
745 unsafe { llama_cpp_sys_2::llama_sampler_chain_n(self.0.sampler.as_ptr()) }
746 }
747
748 /// # Safety
749 ///
750 /// The lifetime of the returned value is tied to `&self`, it must not be
751 /// used beyond that.
752 unsafe fn get(&self, i: c_int) -> Option<ManuallyDrop<LlamaSampler>> {
753 let sampler = unsafe {
754 llama_cpp_sys_2::llama_sampler_chain_get(self.0.sampler.as_mut_ptr_unsound(), i)
755 };
756 // Should only return `None` if out of bounds, we've checked all the
757 // other requirements.
758 let sampler = LlamaSampler::from_raw(sampler)?;
759
760 // HACK: Wrap in `ManuallyDrop`, the lifetime of the sampler is tied
761 // to `&self`.
762 Some(ManuallyDrop::new(sampler))
763 }
764}
765
766#[cfg(test)]
767mod tests {
768 use super::*;
769
770 #[test]
771 fn debug() {
772 let sampler = LlamaSampler::greedy();
773 assert_eq!(r#"LlamaSampler { name: "greedy" }"#, format!("{sampler:?}"));
774
775 let sampler = LlamaSampler::chain_simple([
776 LlamaSampler::dist(1234),
777 LlamaSampler::greedy(),
778 LlamaSampler::chain_simple([LlamaSampler::min_p(1.0, 2)]),
779 ]);
780 assert_eq!(
781 r#"LlamaSampler { name: "chain", chain: [LlamaSampler { name: "dist" }, LlamaSampler { name: "greedy" }, LlamaSampler { name: "chain", chain: [LlamaSampler { name: "min-p" }] }] }"#,
782 format!("{sampler:?}")
783 );
784 }
785}