llama_cpp_bindings/token/
data.rs

1//! Safe wrapper around `llama_token_data`.
2use crate::token::LlamaToken;
3
4/// A transparent wrapper around `llama_token_data`.
5///
6/// Do not rely on `repr(transparent)` for this type. It should be considered an implementation
7/// detail and may change across minor versions.
8#[derive(Clone, Copy, Debug, PartialEq)]
9#[repr(transparent)]
10pub struct LlamaTokenData {
11    data: llama_cpp_bindings_sys::llama_token_data,
12}
13
14impl LlamaTokenData {
15    /// Create a new token data from a token, logit, and probability.
16    /// ```
17    /// # use llama_cpp_bindings::token::LlamaToken;
18    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
19    /// let token = LlamaToken::new(1);
20    /// let token_data = LlamaTokenData::new(token, 1.0, 1.0);
21    #[must_use]
22    pub fn new(LlamaToken(id): LlamaToken, logit: f32, p: f32) -> Self {
23        LlamaTokenData {
24            data: llama_cpp_bindings_sys::llama_token_data { id, logit, p },
25        }
26    }
27    /// Get the token's id
28    /// ```
29    /// # use llama_cpp_bindings::token::LlamaToken;
30    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
31    /// let token = LlamaToken::new(1);
32    /// let token_data = LlamaTokenData::new(token, 1.0, 1.0);
33    /// assert_eq!(token_data.id(), token);
34    /// ```
35    #[must_use]
36    pub fn id(&self) -> LlamaToken {
37        LlamaToken(self.data.id)
38    }
39
40    /// Get the token's logit
41    /// ```
42    /// # use llama_cpp_bindings::token::LlamaToken;
43    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
44    /// let token = LlamaToken::new(1);
45    /// let token_data = LlamaTokenData::new(token, 1.0, 1.0);
46    /// assert_eq!(token_data.logit(), 1.0);
47    /// ```
48    #[must_use]
49    pub fn logit(&self) -> f32 {
50        self.data.logit
51    }
52
53    /// Get the token's probability
54    /// ```
55    /// # use llama_cpp_bindings::token::LlamaToken;
56    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
57    /// let token = LlamaToken::new(1);
58    /// let token_data = LlamaTokenData::new(token, 1.0, 1.0);
59    /// assert_eq!(token_data.p(), 1.0);
60    /// ```
61    #[must_use]
62    pub fn p(&self) -> f32 {
63        self.data.p
64    }
65
66    /// Set the token's id
67    /// ```
68    /// # use llama_cpp_bindings::token::LlamaToken;
69    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
70    /// let token = LlamaToken::new(1);
71    /// let mut token_data = LlamaTokenData::new(token, 1.0, 1.0);
72    /// token_data.set_id(LlamaToken::new(2));
73    /// assert_eq!(token_data.id(), LlamaToken::new(2));
74    /// ```
75    pub fn set_id(&mut self, id: LlamaToken) {
76        self.data.id = id.0;
77    }
78
79    /// Set the token's logit
80    /// ```
81    /// # use llama_cpp_bindings::token::LlamaToken;
82    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
83    /// let token = LlamaToken::new(1);
84    /// let mut token_data = LlamaTokenData::new(token, 1.0, 1.0);
85    /// token_data.set_logit(2.0);
86    /// assert_eq!(token_data.logit(), 2.0);
87    /// ```
88    pub fn set_logit(&mut self, logit: f32) {
89        self.data.logit = logit;
90    }
91
92    /// Set the token's probability
93    /// ```
94    /// # use llama_cpp_bindings::token::LlamaToken;
95    /// # use llama_cpp_bindings::token::data::LlamaTokenData;
96    /// let token = LlamaToken::new(1);
97    /// let mut token_data = LlamaTokenData::new(token, 1.0, 1.0);
98    /// token_data.set_p(2.0);
99    /// assert_eq!(token_data.p(), 2.0);
100    /// ```
101    pub fn set_p(&mut self, p: f32) {
102        self.data.p = p;
103    }
104}
llama_cpp_bindings/token/data.rs

llama_cpp_bindings/token/
data.rs
