Skip to main content

hot_loop/core/
settings.rs

1#[derive(Debug, Clone, Copy)]
2pub struct Settings {
3    pub sample_len: usize,
4    pub temperature: f64,
5    pub top_p: Option<f64>,
6    pub top_k: Option<usize>,
7    pub repeat_penalty: f32,
8    pub repeat_last_n: usize,
9    pub seed: Seed,
10}
11
12#[derive(Debug, Clone, Copy)]
13pub enum Seed {
14    // Random,
15    Custom(u64),
16    Default, // temporary
17}
18
19impl Default for Settings {
20    fn default() -> Self {
21        Self {
22            sample_len: 256,
23            temperature: 0.7,
24            top_p: None,
25            top_k: None,
26            repeat_penalty: 1.1,
27            repeat_last_n: 64,
28            seed: Seed::Default,
29        }
30    }
31}
32
33impl Settings {
34    pub fn with_sample_len(mut self, len: usize) -> Self {
35        self.sample_len = len;
36        self
37    }
38    
39    pub fn with_temperature(mut self, temperature: f64) -> Self {
40        self.temperature = temperature;
41        self
42    }
43    
44    pub fn with_top_p(mut self, top_p: Option<f64>) -> Self {
45        self.top_p = top_p;
46        self
47    }
48    
49    pub fn with_top_k(mut self, top_k: Option<usize>) -> Self {
50        self.top_k = top_k;
51        self
52    }
53    
54    pub fn with_repeat_penalty(mut self, repeat_penalty: f32) -> Self {
55        self.repeat_penalty = repeat_penalty;
56        self
57    }
58    
59    pub fn with_repeat_last_n(mut self, repeat_last_n: usize) -> Self {
60        self.repeat_last_n = repeat_last_n;
61        self
62    }
63    
64    pub fn with_seed(mut self, seed: Seed) -> Self {
65        self.seed = seed;
66        self
67    }
68}