ik_llama_cpp_2/token/
data.rs

1//! Safe wrapper around `llama_token_data`.
2use crate::token::LlamaToken;
3
4use ik_llama_cpp_sys as sys;
5
6/// A transparent wrapper around `llama_token_data` (a token id + its logit and
7/// probability).
8///
9/// Do not rely on `repr(transparent)` for this type — it is an implementation
10/// detail (it lets `[LlamaTokenData]` be reinterpreted as `[llama_token_data]`
11/// for the C sampling calls) and may change.
12#[derive(Clone, Copy, Debug, PartialEq)]
13#[repr(transparent)]
14pub struct LlamaTokenData {
15    data: sys::llama_token_data,
16}
17
18impl LlamaTokenData {
19    /// Create a new token data from a token, logit, and probability.
20    #[must_use]
21    pub fn new(LlamaToken(id): LlamaToken, logit: f32, p: f32) -> Self {
22        LlamaTokenData {
23            data: sys::llama_token_data { id, logit, p },
24        }
25    }
26
27    /// The token's id.
28    #[must_use]
29    pub fn id(&self) -> LlamaToken {
30        LlamaToken(self.data.id)
31    }
32
33    /// The token's logit.
34    #[must_use]
35    pub fn logit(&self) -> f32 {
36        self.data.logit
37    }
38
39    /// The token's probability.
40    #[must_use]
41    pub fn p(&self) -> f32 {
42        self.data.p
43    }
44
45    /// Set the token's id.
46    pub fn set_id(&mut self, id: LlamaToken) {
47        self.data.id = id.0;
48    }
49
50    /// Set the token's logit.
51    pub fn set_logit(&mut self, logit: f32) {
52        self.data.logit = logit;
53    }
54
55    /// Set the token's probability.
56    pub fn set_p(&mut self, p: f32) {
57        self.data.p = p;
58    }
59}
ik_llama_cpp_2/token/data.rs

ik_llama_cpp_2/token/
data.rs
