babichjacob_llama_cpp_2/
sampling.rs

1//! Safe wrapper around `llama_sampler`.
2pub mod params;
3
4use std::ffi::CString;
5use std::fmt::{Debug, Formatter};
6use std::ptr::NonNull;
7
8use crate::context::LlamaContext;
9use crate::model::LlamaModel;
10use crate::token::LlamaToken;
11use crate::LlamaSamplerError;
12
13/// A safe wrapper around `llama_sampler`.
14pub struct LlamaSampler {
15    pub(crate) sampler: NonNull<llama_cpp_sys_2::llama_sampler>,
16}
17
18impl Debug for LlamaSampler {
19    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
20        f.debug_struct("LlamaSamplerChain").finish()
21    }
22}
23
24impl LlamaSampler {
25    /// Create a new `LlamaSampler` from the given parameters.
26    /// # Errors
27    /// Returns an error if the underlying C++ code returns a null pointer.
28    pub fn new(params: params::LlamaSamplerChainParams) -> Result<Self, LlamaSamplerError> {
29        let sampler = unsafe {
30            NonNull::new(llama_cpp_sys_2::llama_sampler_chain_init(
31                params.sampler_chain_params,
32            ))
33            .ok_or(LlamaSamplerError::NullReturn)
34        }?;
35
36        Ok(Self { sampler })
37    }
38
39    /// Samples the token with the largest probability.
40    #[must_use]
41    #[allow(unused_mut)]
42    pub fn add_greedy(mut self) -> Self {
43        unsafe {
44            let greedy_sampler = llama_cpp_sys_2::llama_sampler_init_greedy();
45            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), greedy_sampler);
46        }
47
48        self
49    }
50
51    /// Samples according to the probability distribution of the tokens.
52    #[must_use]
53    #[allow(unused_mut)]
54    pub fn add_dist(mut self, seed: u32) -> Self {
55        unsafe {
56            let dist_sampler = llama_cpp_sys_2::llama_sampler_init_dist(seed);
57            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), dist_sampler);
58        }
59
60        self
61    }
62
63    /// Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" <https://arxiv.org/abs/1904.09751>
64    #[must_use]
65    #[allow(unused_mut)]
66    pub fn add_top_k(mut self, k: i32) -> Self {
67        unsafe {
68            let top_k_sampler = llama_cpp_sys_2::llama_sampler_init_top_k(k);
69            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), top_k_sampler);
70        }
71
72        self
73    }
74
75    /// Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" <https://arxiv.org/abs/1904.09751>
76    #[must_use]
77    #[allow(unused_mut)]
78    pub fn add_top_p(mut self, p: f32, min_keep: usize) -> Self {
79        unsafe {
80            let top_p_sampler = llama_cpp_sys_2::llama_sampler_init_top_p(p, min_keep);
81            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), top_p_sampler);
82        }
83
84        self
85    }
86
87    /// Minimum P sampling as described in <https://github.com/ggerganov/llama.cpp/pull/3841>
88    #[must_use]
89    #[allow(unused_mut)]
90    pub fn add_min_p(mut self, p: f32, min_keep: usize) -> Self {
91        unsafe {
92            let min_p_sampler = llama_cpp_sys_2::llama_sampler_init_min_p(p, min_keep);
93            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), min_p_sampler);
94        }
95
96        self
97    }
98
99    /// Locally Typical Sampling implementation described in the paper <https://arxiv.org/abs/2202.00666>.
100    #[must_use]
101    #[allow(unused_mut)]
102    pub fn add_typical(mut self, p: f32, min_keep: usize) -> Self {
103        unsafe {
104            let typical_sampler = llama_cpp_sys_2::llama_sampler_init_typical(p, min_keep);
105            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), typical_sampler);
106        }
107
108        self
109    }
110
111    /// Updates the logits l_i` = l_i/t. When t <= 0.0f, the maximum logit is kept at it's original value, the rest are set to -inf
112    #[must_use]
113    #[allow(unused_mut)]
114    pub fn add_temp(mut self, t: f32) -> Self {
115        unsafe {
116            let temp_sampler = llama_cpp_sys_2::llama_sampler_init_temp(t);
117            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), temp_sampler);
118        }
119
120        self
121    }
122
123    /// Dynamic temperature implementation (a.k.a. entropy) described in the paper <https://arxiv.org/abs/2309.02772>.
124    #[must_use]
125    #[allow(unused_mut)]
126    pub fn add_temp_ext(mut self, t: f32, delta: f32, exponent: f32) -> Self {
127        unsafe {
128            let temp_ext_sampler = llama_cpp_sys_2::llama_sampler_init_temp_ext(t, delta, exponent);
129            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), temp_ext_sampler);
130        }
131
132        self
133    }
134
135    /// Mirostat 1.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
136    ///
137    /// # Arguments
138    ///
139    /// * `candidates` - A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text.
140    /// * `tau` -  The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text.
141    /// * `eta` - The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates.
142    /// * `m` - The number of tokens considered in the estimation of `s_hat`. This is an arbitrary value that is used to calculate `s_hat`, which in turn helps to calculate the value of `k`. In the paper, they use `m = 100`, but you can experiment with different values to see how it affects the performance of the algorithm.
143    /// * `mu` - Maximum cross-entropy. This value is initialized to be twice the target cross-entropy (`2 * tau`) and is updated in the algorithm based on the error between the target and observed surprisal.
144    #[must_use]
145    #[allow(unused_mut)]
146    pub fn add_mirostat(mut self, n_vocab: i32, seed: u32, tau: f32, eta: f32, m: i32) -> Self {
147        unsafe {
148            let temp_ext_sampler =
149                llama_cpp_sys_2::llama_sampler_init_mirostat(n_vocab, seed, tau, eta, m);
150            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), temp_ext_sampler);
151        }
152
153        self
154    }
155
156    /// Mirostat 2.0 algorithm described in the paper <https://arxiv.org/abs/2007.14966>. Uses tokens instead of words.
157    ///
158    /// # Arguments
159    ///
160    /// * `candidates` - A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text.
161    /// * `tau` -  The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text.
162    /// * `eta` - The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates.
163    /// * `mu` - Maximum cross-entropy. This value is initialized to be twice the target cross-entropy (`2 * tau`) and is updated in the algorithm based on the error between the target and observed surprisal.
164    #[must_use]
165    #[allow(unused_mut)]
166    pub fn add_mirostat_v2(mut self, seed: u32, tau: f32, eta: f32) -> Self {
167        unsafe {
168            let temp_ext_sampler = llama_cpp_sys_2::llama_sampler_init_mirostat_v2(seed, tau, eta);
169            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), temp_ext_sampler);
170        }
171
172        self
173    }
174
175    /// Samples constrained by a context-free grammar in the GGML BNF (GBNF) format.
176    ///
177    /// # Panics
178    /// Panics if a provided string contains a null byte.
179    #[must_use]
180    #[allow(unused_mut)]
181    pub fn add_grammar(
182        mut self,
183        model: &LlamaModel,
184        grammar_str: &str,
185        grammar_root: &str,
186    ) -> Self {
187        unsafe {
188            let grammar_str = CString::new(grammar_str).unwrap();
189            let grammar_root = CString::new(grammar_root).unwrap();
190            let grammar_sampler = llama_cpp_sys_2::llama_sampler_init_grammar(
191                model.model.as_ptr(),
192                grammar_str.as_ptr(),
193                grammar_root.as_ptr(),
194            );
195            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), grammar_sampler);
196        }
197
198        self
199    }
200
201    /// Adds penalties to the sampler. This can be used to penalize certain patterns in the generated text, such as repeating the same token multiple times or using the same token too frequently.
202    #[allow(unused_mut, clippy::too_many_arguments)]
203    #[must_use]
204    pub fn add_penalties(
205        mut self,
206        n_vocab: i32,
207        special_eos_id: i32,
208        linefeed_id: i32,
209        penalty_last_n: i32,
210        penalty_repeat: f32,
211        penalty_freq: f32,
212        penalty_present: f32,
213        penalize_nl: bool,
214        ignore_eos: bool,
215    ) -> Self {
216        unsafe {
217            let temp_ext_sampler = llama_cpp_sys_2::llama_sampler_init_penalties(
218                n_vocab,
219                special_eos_id,
220                linefeed_id,
221                penalty_last_n,
222                penalty_repeat,
223                penalty_freq,
224                penalty_present,
225                penalize_nl,
226                ignore_eos,
227            );
228            llama_cpp_sys_2::llama_sampler_chain_add(self.sampler.as_ptr(), temp_ext_sampler);
229        }
230
231        self
232    }
233
234    /// Sample and accept a token from the idx-th output of the last evaluation
235    #[must_use]
236    pub fn sample(&self, ctx: &LlamaContext, idx: i32) -> LlamaToken {
237        let token = unsafe {
238            llama_cpp_sys_2::llama_sampler_sample(self.sampler.as_ptr(), ctx.context.as_ptr(), idx)
239        };
240
241        LlamaToken(token)
242    }
243
244    /// Accepts a token from the sampler, possibly updating the internal state of certain samplers (e.g. grammar, repetition, etc.)
245    pub fn accept(&mut self, token: LlamaToken) {
246        unsafe { llama_cpp_sys_2::llama_sampler_accept(self.sampler.as_ptr(), token.0) }
247    }
248}
249
250impl Drop for LlamaSampler {
251    fn drop(&mut self) {
252        unsafe {
253            llama_cpp_sys_2::llama_sampler_free(self.sampler.as_ptr());
254        }
255    }
256}
babichjacob_llama_cpp_2/sampling.rs

babichjacob_llama_cpp_2/
sampling.rs
