litertlm 0.2.1

Safe, idiomatic Rust bindings to Google's LiteRT-LM on-device LLM engine.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

//! Sampling parameters for text generation.

use litert_lm_sys as sys;

/// The sampling strategy used during token selection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Sampler {
    /// Pick the token with maximum logit (argmax).
    Greedy,
    /// Probabilistically pick among the top-k tokens.
    TopK,
    /// Probabilistically pick from the nucleus (top-p).
    TopP,
}

impl Sampler {
    fn to_raw(self) -> sys::Type {
        match self {
            Self::Greedy => sys::kGreedy,
            Self::TopK => sys::kTopK,
            Self::TopP => sys::kTopP,
        }
    }
}

/// Parameters controlling how the model selects the next token.
///
/// # Example
///
/// ```no_run
/// use litertlm::SamplerParams;
///
/// let params = SamplerParams::default()
///     .top_k(40)
///     .top_p(0.95)
///     .temperature(0.8)
///     .seed(42);
/// ```
#[derive(Debug, Clone)]
pub struct SamplerParams {
    pub(crate) sampler: Sampler,
    pub(crate) top_k: i32,
    pub(crate) top_p: f32,
    pub(crate) temperature: f32,
    pub(crate) seed: i32,
}

impl Default for SamplerParams {
    fn default() -> Self {
        Self {
            sampler: Sampler::TopK,
            top_k: 40,
            top_p: 0.95,
            temperature: 0.8,
            seed: 0,
        }
    }
}

impl SamplerParams {
    /// Set the top-k value. Implies [`Sampler::TopK`] if not already set.
    #[must_use]
    pub fn top_k(mut self, v: i32) -> Self {
        self.top_k = v;
        self.sampler = Sampler::TopK;
        self
    }

    /// Set the nucleus probability threshold. Implies [`Sampler::TopP`].
    #[must_use]
    pub fn top_p(mut self, v: f32) -> Self {
        self.top_p = v;
        self.sampler = Sampler::TopP;
        self
    }

    /// Set the softmax temperature.
    #[must_use]
    pub fn temperature(mut self, v: f32) -> Self {
        self.temperature = v;
        self
    }

    /// Set the random seed for reproducible generation.
    #[must_use]
    pub fn seed(mut self, v: i32) -> Self {
        self.seed = v;
        self
    }

    /// Use greedy decoding (argmax, deterministic).
    #[must_use]
    pub fn greedy(mut self) -> Self {
        self.sampler = Sampler::Greedy;
        self
    }

    pub(crate) fn to_raw(&self) -> sys::LiteRtLmSamplerParams {
        sys::LiteRtLmSamplerParams {
            type_: self.sampler.to_raw(),
            top_k: self.top_k,
            top_p: self.top_p,
            temperature: self.temperature,
            seed: self.seed,
        }
    }
}