ik_llama_cpp_2/token/
data_array.rs

1//! A rusty equivalent of `llama_token_data_array`.
2use std::ptr;
3
4use ik_llama_cpp_sys as sys;
5
6use crate::context::LlamaContext;
7use crate::token::data::LlamaTokenData;
8use crate::token::LlamaToken;
9
10/// A safe wrapper around `llama_token_data_array` — the candidate set consumed
11/// (and mutated) by sampling.
12#[derive(Debug, Clone, PartialEq)]
13pub struct LlamaTokenDataArray {
14    /// The underlying candidate data.
15    pub data: Vec<LlamaTokenData>,
16    /// The index of the selected token in `data`, if a sampler has chosen one.
17    pub selected: Option<usize>,
18    /// Whether `data` is sorted by logit (descending).
19    pub sorted: bool,
20}
21
22impl LlamaTokenDataArray {
23    /// Create a new `LlamaTokenDataArray` from a vector and whether the data is
24    /// sorted by logit.
25    #[must_use]
26    pub fn new(data: Vec<LlamaTokenData>, sorted: bool) -> Self {
27        Self {
28            data,
29            selected: None,
30            sorted,
31        }
32    }
33
34    /// Create a new `LlamaTokenDataArray` from an iterator of [`LlamaTokenData`].
35    pub fn from_iter<T>(data: T, sorted: bool) -> LlamaTokenDataArray
36    where
37        T: IntoIterator<Item = LlamaTokenData>,
38    {
39        Self::new(data.into_iter().collect(), sorted)
40    }
41
42    /// Build from a logits slice (`id` = index). Probabilities start at 0 and
43    /// the data is unsorted.
44    #[must_use]
45    pub fn from_logits(logits: &[f32]) -> Self {
46        let data = logits
47            .iter()
48            .enumerate()
49            .map(|(i, &logit)| LlamaTokenData::new(LlamaToken(i as sys::llama_token), logit, 0.0))
50            .collect();
51        Self::new(data, false)
52    }
53
54    /// The currently selected token, if a sampler has chosen one.
55    #[must_use]
56    pub fn selected_token(&self) -> Option<LlamaToken> {
57        self.data.get(self.selected?).map(LlamaTokenData::id)
58    }
59
60    /// Apply repetition / frequency / presence penalties in place over the most
61    /// recent `last_tokens`.
62    ///
63    /// Wraps ik's legacy `llama_sample_repetition_penalties`. An empty
64    /// `last_tokens` slice is a no-op. Logits are mutated in place; candidate
65    /// order is preserved.
66    pub fn apply_repetition_penalties(
67        &mut self,
68        ctx: &mut LlamaContext,
69        last_tokens: &[LlamaToken],
70        penalty_repeat: f32,
71        penalty_freq: f32,
72        penalty_present: f32,
73    ) {
74        let raw: Vec<sys::llama_token> = last_tokens.iter().map(|t| t.0).collect();
75        let ctx_ptr = ctx.as_ptr();
76        // SAFETY: valid ctx; `c` describes `self.data` (mutated in place through
77        // its pointer); `raw` lives for the call.
78        unsafe {
79            self.modify_as_c_llama_token_data_array(|c| {
80                sys::llama_sample_repetition_penalties(
81                    ctx_ptr,
82                    c,
83                    raw.as_ptr(),
84                    raw.len(),
85                    penalty_repeat,
86                    penalty_freq,
87                    penalty_present,
88                );
89            });
90        }
91    }
92
93    /// Build a `llama_token_data_array` snapshot pointing at `self.data`.
94    ///
95    /// The result borrows `self.data`'s buffer; it is valid only while `self` is
96    /// not moved/reallocated. C sampling calls that only rewrite logits in place
97    /// (no resize) may use this directly; anything that can resize / reselect
98    /// must go through [`Self::modify_as_c_llama_token_data_array`].
99    pub(crate) fn as_c(&mut self) -> sys::llama_token_data_array {
100        sys::llama_token_data_array {
101            data: self.data.as_mut_ptr().cast::<sys::llama_token_data>(),
102            size: self.data.len(),
103            selected: self.selected.and_then(|s| s.try_into().ok()).unwrap_or(-1),
104            sorted: self.sorted,
105        }
106    }
107
108    /// Expose `self` as a C `llama_token_data_array` for the duration of
109    /// `modify`, then sync `size` / `sorted` / `selected` back.
110    ///
111    /// # Safety
112    ///
113    /// `modify` must leave the array's `data`/`size` describing initialized
114    /// token data within this buffer's capacity, and must set `sorted` honestly.
115    pub(crate) unsafe fn modify_as_c_llama_token_data_array<T>(
116        &mut self,
117        modify: impl FnOnce(&mut sys::llama_token_data_array) -> T,
118    ) -> T {
119        let data = self.data.as_mut_ptr().cast::<sys::llama_token_data>();
120        let mut c = sys::llama_token_data_array {
121            data,
122            size: self.data.len(),
123            selected: self.selected.and_then(|s| s.try_into().ok()).unwrap_or(-1),
124            sorted: self.sorted,
125        };
126
127        let result = modify(&mut c);
128
129        assert!(
130            c.size <= self.data.capacity(),
131            "sampler returned an array larger than the candidate buffer capacity"
132        );
133        if !ptr::eq(c.data, data) {
134            ptr::copy(c.data, data, c.size);
135        }
136        // SAFETY: `modify` guarantees `0..c.size` is initialized (see contract).
137        self.data.set_len(c.size);
138        self.sorted = c.sorted;
139        self.selected = c.selected.try_into().ok().filter(|&s| s < self.data.len());
140        result
141    }
142}
ik_llama_cpp_2/token/data_array.rs

ik_llama_cpp_2/token/
data_array.rs
