kalosm_language_model/model/
generation_parameters.rs

1#[cfg(feature = "sample")]
2use std::hash::Hash;
3#[cfg(feature = "sample")]
4use std::hash::Hasher;
5
6#[cfg(feature = "sample")]
7use llm_samplers::configure::SamplerChainBuilder;
8#[cfg(feature = "sample")]
9use llm_samplers::prelude::*;
10
11/// Parameters to use when generating text.
12#[derive(Debug)]
13pub struct GenerationParameters {
14    pub(crate) temperature: f32,
15    pub(crate) tau: f32,
16    pub(crate) eta: f32,
17    pub(crate) mu: f32,
18    pub(crate) top_p: f64,
19    pub(crate) top_k: u32,
20    pub(crate) repetition_penalty: f32,
21    pub(crate) repetition_penalty_range: u32,
22    pub(crate) max_length: u32,
23    pub(crate) stop_on: Option<String>,
24    pub(crate) seed: Option<u64>,
25    #[cfg(feature = "sample")]
26    sampler: Option<(u64, SamplerChain)>,
27}
28
29impl PartialEq for GenerationParameters {
30    fn eq(&self, other: &Self) -> bool {
31        self.temperature == other.temperature
32            && self.eta == other.eta
33            && self.tau == other.tau
34            && self.mu == other.mu
35            && self.top_p == other.top_p
36            && self.repetition_penalty == other.repetition_penalty
37            && self.repetition_penalty_range == other.repetition_penalty_range
38            && self.max_length == other.max_length
39            && self.stop_on == other.stop_on
40    }
41}
42
43impl Clone for GenerationParameters {
44    fn clone(&self) -> Self {
45        Self {
46            temperature: self.temperature,
47            eta: self.eta,
48            tau: self.tau,
49            mu: self.mu,
50            top_p: self.top_p,
51            top_k: self.top_k,
52            repetition_penalty: self.repetition_penalty,
53            repetition_penalty_range: self.repetition_penalty_range,
54            max_length: self.max_length,
55            stop_on: self.stop_on.clone(),
56            seed: None,
57            #[cfg(feature = "sample")]
58            sampler: None,
59        }
60    }
61}
62
63impl Default for GenerationParameters {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69#[cfg(feature = "sample")]
70impl Sampler for GenerationParameters {
71    fn sample<'a>(
72        &mut self,
73        res: &mut dyn HasSamplerResources,
74        logits: &'a mut Logits,
75    ) -> anyhow::Result<&'a mut Logits> {
76        self.with_sampler(|sampler| sampler.sample(res, logits))
77    }
78
79    fn sample_token(
80        &mut self,
81        res: &mut dyn HasSamplerResources,
82        logits: &mut Logits,
83    ) -> anyhow::Result<Option<TID>> {
84        self.with_sampler(|sampler| sampler.sample_token(res, logits))
85    }
86
87    fn sampled_token_id(&self) -> Option<TID> {
88        self.sampler().sampled_token_id()
89    }
90}
91
92impl GenerationParameters {
93    /// Create a new [`GenerationParameters`]
94    pub const fn new() -> Self {
95        Self {
96            temperature: 0.8,
97            eta: 0.1,
98            tau: 5.,
99            mu: 10.,
100            top_p: 1.0,
101            top_k: 1,
102            repetition_penalty: 1.3,
103            repetition_penalty_range: 64,
104            max_length: u32::MAX,
105            stop_on: None,
106            seed: None,
107            #[cfg(feature = "sample")]
108            sampler: None,
109        }
110    }
111
112    #[cfg(feature = "sample")]
113    fn with_sampler<O>(&mut self, with_sampler: impl FnOnce(&mut SamplerChain) -> O) -> O {
114        let mut hash = std::collections::hash_map::DefaultHasher::new();
115        self.eta.to_le_bytes().hash(&mut hash);
116        self.mu.to_le_bytes().hash(&mut hash);
117        self.repetition_penalty.to_le_bytes().hash(&mut hash);
118        self.repetition_penalty_range.hash(&mut hash);
119        self.tau.to_le_bytes().hash(&mut hash);
120        self.top_p.to_le_bytes().hash(&mut hash);
121        self.temperature.to_le_bytes().hash(&mut hash);
122        self.max_length.hash(&mut hash);
123        let hash = hash.finish();
124        if let Some((old_hash, sampler)) = &mut self.sampler {
125            if *old_hash == hash {
126                return with_sampler(sampler);
127            }
128        }
129        let mut sampler = self.sampler();
130        let output = with_sampler(&mut sampler);
131        self.sampler = Some((hash, sampler));
132        output
133    }
134
135    #[cfg(feature = "sample")]
136    /// Create a sampler chain from the generation parameters.
137    pub fn sampler(&self) -> SamplerChain {
138        use llm_samplers::configure::SamplerSlot;
139        let GenerationParameters {
140            temperature,
141            tau,
142            eta,
143            mu,
144            repetition_penalty,
145            repetition_penalty_range,
146            top_p: _,
147            max_length: _,
148            stop_on: _,
149            ..
150        } = self;
151        let temperature = *temperature;
152        let tau = *tau;
153        let eta = *eta;
154        let mu = *mu;
155        let repetition_penalty = *repetition_penalty;
156        let repetition_penalty_range = *repetition_penalty_range;
157        SamplerChainBuilder::from([
158            (
159                "repetition",
160                SamplerSlot::new_static(move || {
161                    Box::new(
162                        SampleRepetition::default()
163                            .penalty(repetition_penalty)
164                            .last_n(repetition_penalty_range as usize),
165                    )
166                }),
167            ),
168            (
169                "freqpresence",
170                SamplerSlot::new_static(move || Box::new(SampleFreqPresence::default().last_n(64))),
171            ),
172            (
173                "seqrepetition",
174                SamplerSlot::new_static(move || Box::<SampleSeqRepetition>::default()),
175            ),
176            (
177                "temperature",
178                SamplerSlot::new_static(move || {
179                    Box::new(SampleTemperature::default().temperature(temperature))
180                }),
181            ),
182            (
183                "mirostat2",
184                SamplerSlot::new_static(move || {
185                    Box::new(SampleMirostat2::default().tau(tau).eta(eta).mu(mu))
186                }),
187            ),
188        ])
189        .into_chain()
190    }
191
192    /// Set the top_p parameter to the generation parameters (only used by the OpenAI API).
193    pub fn with_top_p(mut self, top_p: f64) -> Self {
194        self.top_p = top_p;
195        self
196    }
197
198    /// Set the top_k parameter to the generation parameters (only used by the Anthropic API).
199    pub fn with_top_k(mut self, top_k: u32) -> Self {
200        self.top_k = top_k;
201        self
202    }
203
204    #[cfg(feature = "sample")]
205    /// Get the mirostat2 sampler from the generation parameters.
206    pub fn mirostat2_sampler(self) -> SampleMirostat2 {
207        SampleMirostat2::default()
208            .tau(self.tau)
209            .eta(self.eta)
210            .mu(self.mu)
211    }
212
213    #[cfg(feature = "sample")]
214    /// Create a sampler chain from the generation parameters without removing any tokens. This can be useful in combination with [`ModelExt::stream_structured_text_with_sampler`] which may pick unlikely tokens.
215    pub fn bias_only_sampler(self) -> SamplerChain {
216        use llm_samplers::configure::SamplerSlot;
217        let GenerationParameters {
218            temperature,
219            repetition_penalty,
220            repetition_penalty_range,
221            ..
222        } = self;
223        SamplerChainBuilder::from([
224            (
225                "repetition",
226                SamplerSlot::new_static(move || {
227                    Box::new(
228                        SampleRepetition::default()
229                            .penalty(repetition_penalty)
230                            .last_n(repetition_penalty_range as usize),
231                    )
232                }),
233            ),
234            (
235                "freqpresence",
236                SamplerSlot::new_static(move || Box::new(SampleFreqPresence::default().last_n(64))),
237            ),
238            (
239                "seqrepetition",
240                SamplerSlot::new_static(move || Box::<SampleSeqRepetition>::default()),
241            ),
242            (
243                "temperature",
244                SamplerSlot::new_static(move || {
245                    Box::new(SampleTemperature::default().temperature(temperature))
246                }),
247            ),
248        ])
249        .into_chain()
250    }
251
252    /// Set the temperature to use when generating text.
253    pub fn with_temperature(mut self, temperature: f32) -> Self {
254        self.temperature = temperature;
255        self
256    }
257
258    /// Set the tau to use when generating text.
259    pub fn with_tau(mut self, tau: f32) -> Self {
260        self.tau = tau;
261        self
262    }
263
264    /// Set the eta to use when generating text.
265    pub fn with_eta(mut self, eta: f32) -> Self {
266        self.eta = eta;
267        self
268    }
269
270    /// Set the mu to use when generating text.
271    pub fn with_mu(mut self, mu: f32) -> Self {
272        self.mu = mu;
273        self
274    }
275
276    /// Set the repetition penalty to use when generating text.
277    pub fn with_repetition_penalty(mut self, repetition_penalty: f32) -> Self {
278        self.repetition_penalty = repetition_penalty;
279        self
280    }
281
282    /// Set the repetition penalty range to use when generating text.
283    pub fn with_repetition_penalty_range(mut self, repetition_penalty_range: u32) -> Self {
284        self.repetition_penalty_range = repetition_penalty_range;
285        self
286    }
287
288    /// Set the maximum length to use when generating text.
289    pub fn with_max_length(mut self, max_length: u32) -> Self {
290        self.max_length = max_length;
291        self
292    }
293
294    /// Set the string to stop on when generating text.
295    pub fn with_stop_on(mut self, stop_on: impl Into<Option<String>>) -> Self {
296        self.stop_on = stop_on.into();
297        self
298    }
299
300    /// Set the seed to use when generating text.
301    pub fn with_seed(mut self, seed: impl Into<Option<u64>>) -> Self {
302        self.seed = seed.into();
303        self
304    }
305
306    /// Get the temperature to use when generating text.
307    pub fn temperature(&self) -> f32 {
308        self.temperature
309    }
310
311    /// Get the tau to use when generating text.
312    pub fn tau(&self) -> f32 {
313        self.tau
314    }
315
316    /// Get the eta to use when generating text.
317    pub fn eta(&self) -> f32 {
318        self.eta
319    }
320
321    /// Get the mu to use when generating text.
322    pub fn mu(&self) -> f32 {
323        self.mu
324    }
325
326    /// Get the repetition penalty to use when generating text.
327    pub fn repetition_penalty(&self) -> f32 {
328        self.repetition_penalty
329    }
330
331    /// Get the repetition penalty range to use when generating text.
332    pub fn repetition_penalty_range(&self) -> u32 {
333        self.repetition_penalty_range
334    }
335
336    /// Get the maximum length to use when generating text.
337    pub fn max_length(&self) -> u32 {
338        self.max_length
339    }
340
341    /// Get the string to stop on when generating text.
342    pub fn stop_on(&self) -> Option<&str> {
343        self.stop_on.as_deref()
344    }
345
346    /// Get the seed to use when generating text.
347    pub fn seed(&self) -> Option<u64> {
348        self.seed
349    }
350}
kalosm_language_model/model/generation_parameters.rs

kalosm_language_model/model/
generation_parameters.rs
