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