llama_cpp_bindings/token/
data_array.rs

1use std::ptr;
2
3use crate::error::SamplerApplyError;
4use crate::error::TokenSamplingError;
5use crate::sampling::LlamaSampler;
6use crate::token::data::LlamaTokenData;
7
8use super::LlamaToken;
9
10fn sampler_apply_status_to_result(
11    status: llama_cpp_bindings_sys::llama_rs_sampler_apply_status,
12    out_error: *mut std::os::raw::c_char,
13) -> Result<(), SamplerApplyError> {
14    match status {
15        llama_cpp_bindings_sys::LLAMA_RS_SAMPLER_APPLY_OK => Ok(()),
16        llama_cpp_bindings_sys::LLAMA_RS_SAMPLER_APPLY_NULL_SAMPLER_ARG => {
17            Err(SamplerApplyError::NullSampler)
18        }
19        llama_cpp_bindings_sys::LLAMA_RS_SAMPLER_APPLY_ERROR_STRING_ALLOCATION_FAILED => {
20            Err(SamplerApplyError::NotEnoughMemory)
21        }
22        llama_cpp_bindings_sys::LLAMA_RS_SAMPLER_APPLY_VENDORED_THREW_CXX_EXCEPTION => {
23            let message = unsafe { crate::ffi_error_reader::read_and_free_cpp_error(out_error) };
24            Err(SamplerApplyError::Reported { message })
25        }
26        other => {
27            unreachable!("llama_rs_sampler_apply returned unrecognized status {other}")
28        }
29    }
30}
31
32#[derive(Debug, Clone, PartialEq)]
33pub struct LlamaTokenDataArray {
34    pub data: Vec<LlamaTokenData>,
35    pub selected: Option<usize>,
36    pub sorted: bool,
37}
38
39impl LlamaTokenDataArray {
40    #[must_use]
41    pub const fn new(data: Vec<LlamaTokenData>, sorted: bool) -> Self {
42        Self {
43            data,
44            selected: None,
45            sorted,
46        }
47    }
48
49    pub fn from_iter<TIterator>(data: TIterator, sorted: bool) -> Self
50    where
51        TIterator: IntoIterator<Item = LlamaTokenData>,
52    {
53        Self::new(data.into_iter().collect(), sorted)
54    }
55
56    #[must_use]
57    pub fn selected_token(&self) -> Option<LlamaToken> {
58        self.data.get(self.selected?).map(LlamaTokenData::id)
59    }
60}
61
62impl LlamaTokenDataArray {
63    /// # Panics
64    ///
65    /// Panics if some of the safety conditions are not met. (we cannot check all of them at
66    /// runtime so breaking them is UB)
67    ///
68    /// # Safety
69    ///
70    /// The returned array formed by the data pointer and the length must entirely consist of
71    /// initialized token data and the length must be less than the capacity of this array's data
72    /// buffer.
73    /// If the data is not sorted, sorted must be false.
74    pub unsafe fn modify_as_c_llama_token_data_array<TResult>(
75        &mut self,
76        modify: impl FnOnce(&mut llama_cpp_bindings_sys::llama_token_data_array) -> TResult,
77    ) -> TResult {
78        let size = self.data.len();
79        let data = self
80            .data
81            .as_mut_ptr()
82            .cast::<llama_cpp_bindings_sys::llama_token_data>();
83
84        let mut c_llama_token_data_array = llama_cpp_bindings_sys::llama_token_data_array {
85            data,
86            size,
87            selected: self
88                .selected
89                .and_then(|selected_index| selected_index.try_into().ok())
90                .unwrap_or(-1),
91            sorted: self.sorted,
92        };
93
94        let result = modify(&mut c_llama_token_data_array);
95
96        assert!(c_llama_token_data_array.size <= self.data.capacity());
97        // SAFETY: caller guarantees the returned data and size are valid.
98        unsafe {
99            if !ptr::eq(c_llama_token_data_array.data, data) {
100                ptr::copy(
101                    c_llama_token_data_array.data,
102                    data,
103                    c_llama_token_data_array.size,
104                );
105            }
106            self.data.set_len(c_llama_token_data_array.size);
107        }
108
109        self.sorted = c_llama_token_data_array.sorted;
110        self.selected = c_llama_token_data_array
111            .selected
112            .try_into()
113            .ok()
114            .filter(|&s| s < self.data.len());
115
116        result
117    }
118
119    /// # Errors
120    ///
121    /// Returns [`SamplerApplyError`] if the sampler pointer is null, the vendored
122    /// sampler runs out of memory, or it throws a C++ exception while applying.
123    pub fn apply_sampler(&mut self, sampler: &LlamaSampler) -> Result<(), SamplerApplyError> {
124        unsafe {
125            self.modify_as_c_llama_token_data_array(|c_llama_token_data_array| {
126                let mut out_error: *mut std::os::raw::c_char = ptr::null_mut();
127                let status = llama_cpp_bindings_sys::llama_rs_sampler_apply(
128                    sampler.sampler,
129                    c_llama_token_data_array,
130                    &raw mut out_error,
131                );
132                sampler_apply_status_to_result(status, out_error)
133            })
134        }
135    }
136
137    /// # Errors
138    /// Returns [`SamplerApplyError`] if applying the sampler fails.
139    pub fn with_sampler(mut self, sampler: &mut LlamaSampler) -> Result<Self, SamplerApplyError> {
140        self.apply_sampler(sampler)?;
141        Ok(self)
142    }
143
144    /// # Errors
145    /// Returns [`TokenSamplingError::SamplerApply`] if applying the sampler fails, or
146    /// [`TokenSamplingError::NoTokenSelected`] if the sampler fails to select a token.
147    pub fn sample_token(&mut self, seed: u32) -> Result<LlamaToken, TokenSamplingError> {
148        self.apply_sampler(&LlamaSampler::dist(seed))?;
149        self.selected_token()
150            .ok_or(TokenSamplingError::NoTokenSelected)
151    }
152
153    /// # Errors
154    /// Returns [`TokenSamplingError::SamplerApply`] if applying the sampler fails, or
155    /// [`TokenSamplingError::NoTokenSelected`] if the sampler fails to select a token.
156    pub fn sample_token_greedy(&mut self) -> Result<LlamaToken, TokenSamplingError> {
157        self.apply_sampler(&LlamaSampler::greedy())?;
158        self.selected_token()
159            .ok_or(TokenSamplingError::NoTokenSelected)
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use crate::error::SamplerApplyError;
166    use crate::token::LlamaToken;
167    use crate::token::data::LlamaTokenData;
168
169    use super::LlamaTokenDataArray;
170    use super::sampler_apply_status_to_result;
171
172    #[test]
173    fn sampler_apply_status_allocation_failed_returns_not_enough_memory() {
174        assert_eq!(
175            sampler_apply_status_to_result(
176                llama_cpp_bindings_sys::LLAMA_RS_SAMPLER_APPLY_ERROR_STRING_ALLOCATION_FAILED,
177                std::ptr::null_mut(),
178            ),
179            Err(SamplerApplyError::NotEnoughMemory),
180        );
181    }
182
183    #[test]
184    fn sampler_apply_status_cxx_exception_returns_reported_with_unknown_message() {
185        assert_eq!(
186            sampler_apply_status_to_result(
187                llama_cpp_bindings_sys::LLAMA_RS_SAMPLER_APPLY_VENDORED_THREW_CXX_EXCEPTION,
188                std::ptr::null_mut(),
189            ),
190            Err(SamplerApplyError::Reported {
191                message: "unknown error".to_owned(),
192            }),
193        );
194    }
195
196    #[test]
197    #[should_panic(expected = "llama_rs_sampler_apply returned unrecognized status")]
198    fn sampler_apply_status_unrecognized_panics() {
199        let _ = sampler_apply_status_to_result(
200            llama_cpp_bindings_sys::llama_rs_sampler_apply_status::MAX,
201            std::ptr::null_mut(),
202        );
203    }
204
205    #[test]
206    fn apply_greedy_sampler_selects_highest_logit() {
207        use crate::sampling::LlamaSampler;
208
209        let mut array = LlamaTokenDataArray::new(
210            vec![
211                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
212                LlamaTokenData::new(LlamaToken::new(1), 5.0, 0.0),
213                LlamaTokenData::new(LlamaToken::new(2), 3.0, 0.0),
214            ],
215            false,
216        );
217
218        array
219            .apply_sampler(&LlamaSampler::greedy())
220            .expect("test: greedy sampler must apply");
221
222        assert_eq!(array.selected_token(), Some(LlamaToken::new(1)));
223    }
224
225    #[test]
226    fn with_sampler_builder_pattern() {
227        use crate::sampling::LlamaSampler;
228
229        let array = LlamaTokenDataArray::new(
230            vec![
231                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
232                LlamaTokenData::new(LlamaToken::new(1), 5.0, 0.0),
233            ],
234            false,
235        )
236        .with_sampler(&mut LlamaSampler::greedy())
237        .expect("test: building with greedy sampler must succeed");
238
239        assert_eq!(array.selected_token(), Some(LlamaToken::new(1)));
240    }
241
242    #[test]
243    fn with_sampler_with_null_sampler_returns_sampler_apply_error() {
244        use crate::sampling::LlamaSampler;
245
246        let mut null_sampler = LlamaSampler {
247            sampler: std::ptr::null_mut(),
248        };
249        let array = LlamaTokenDataArray::new(
250            vec![LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0)],
251            false,
252        );
253
254        assert_eq!(
255            array.with_sampler(&mut null_sampler),
256            Err(SamplerApplyError::NullSampler),
257        );
258    }
259
260    #[test]
261    fn sample_token_greedy_returns_highest() {
262        let mut array = LlamaTokenDataArray::new(
263            vec![
264                LlamaTokenData::new(LlamaToken::new(10), 0.1, 0.0),
265                LlamaTokenData::new(LlamaToken::new(20), 9.9, 0.0),
266            ],
267            false,
268        );
269
270        let token = array
271            .sample_token_greedy()
272            .expect("test: greedy sampler should select a token");
273
274        assert_eq!(token, LlamaToken::new(20));
275    }
276
277    #[test]
278    fn from_iter_creates_array_from_iterator() {
279        let array = LlamaTokenDataArray::from_iter(
280            [
281                LlamaTokenData::new(LlamaToken::new(0), 0.0, 0.0),
282                LlamaTokenData::new(LlamaToken::new(1), 1.0, 0.0),
283                LlamaTokenData::new(LlamaToken::new(2), 2.0, 0.0),
284            ],
285            false,
286        );
287
288        assert_eq!(array.data.len(), 3);
289        assert!(!array.sorted);
290        assert!(array.selected.is_none());
291    }
292
293    #[test]
294    fn sample_token_with_seed_selects_a_token() {
295        let mut array = LlamaTokenDataArray::new(
296            vec![
297                LlamaTokenData::new(LlamaToken::new(10), 1.0, 0.0),
298                LlamaTokenData::new(LlamaToken::new(20), 1.0, 0.0),
299            ],
300            false,
301        );
302
303        let token = array
304            .sample_token(42)
305            .expect("test: dist sampler should select a token");
306
307        assert!(token == LlamaToken::new(10) || token == LlamaToken::new(20));
308    }
309
310    #[test]
311    fn selected_token_returns_none_when_no_selection() {
312        let array = LlamaTokenDataArray::new(
313            vec![LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0)],
314            false,
315        );
316
317        assert!(array.selected_token().is_none());
318    }
319
320    #[test]
321    fn selected_token_returns_none_when_index_out_of_bounds() {
322        let array = LlamaTokenDataArray {
323            data: vec![LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0)],
324            selected: Some(5),
325            sorted: false,
326        };
327
328        assert!(array.selected_token().is_none());
329    }
330
331    #[test]
332    fn modify_as_c_llama_token_data_array_copies_when_data_pointer_changes() {
333        let mut array = LlamaTokenDataArray::new(
334            vec![
335                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
336                LlamaTokenData::new(LlamaToken::new(1), 2.0, 0.0),
337                LlamaTokenData::new(LlamaToken::new(2), 3.0, 0.0),
338            ],
339            false,
340        );
341
342        let replacement = [
343            llama_cpp_bindings_sys::llama_token_data {
344                id: 10,
345                logit: 5.0,
346                p: 0.0,
347            },
348            llama_cpp_bindings_sys::llama_token_data {
349                id: 20,
350                logit: 6.0,
351                p: 0.0,
352            },
353        ];
354
355        unsafe {
356            array.modify_as_c_llama_token_data_array(|c_array| {
357                c_array.data = replacement.as_ptr().cast_mut();
358                c_array.size = replacement.len();
359                c_array.selected = 0;
360            });
361        }
362
363        assert_eq!(array.data.len(), 2);
364        assert_eq!(array.data[0].id(), LlamaToken::new(10));
365        assert_eq!(array.data[1].id(), LlamaToken::new(20));
366        assert_eq!(array.selected, Some(0));
367    }
368
369    #[test]
370    fn apply_sampler_with_null_sampler_returns_null_sampler_error() {
371        use crate::sampling::LlamaSampler;
372
373        let mut array = LlamaTokenDataArray::new(
374            vec![LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0)],
375            false,
376        );
377
378        let null_sampler = LlamaSampler {
379            sampler: std::ptr::null_mut(),
380        };
381
382        assert_eq!(
383            array.apply_sampler(&null_sampler),
384            Err(SamplerApplyError::NullSampler)
385        );
386    }
387
388    #[test]
389    fn modify_clears_selection_when_index_is_out_of_range() {
390        let mut array = LlamaTokenDataArray::new(
391            vec![LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0)],
392            false,
393        );
394
395        unsafe {
396            array.modify_as_c_llama_token_data_array(|c_array| {
397                c_array.selected = 5;
398            });
399        }
400
401        assert_eq!(array.selected, None);
402    }
403
404    #[test]
405    fn selected_overflow_uses_negative_one() {
406        let mut array = LlamaTokenDataArray {
407            data: vec![LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0)],
408            selected: Some(usize::MAX),
409            sorted: false,
410        };
411
412        unsafe {
413            array.modify_as_c_llama_token_data_array(|c_array| {
414                assert_eq!(c_array.selected, -1);
415            });
416        }
417    }
418
419    #[test]
420    fn preset_valid_selection_is_passed_through_as_index() {
421        let mut array = LlamaTokenDataArray {
422            data: vec![
423                LlamaTokenData::new(LlamaToken::new(0), 1.0, 0.0),
424                LlamaTokenData::new(LlamaToken::new(1), 2.0, 0.0),
425            ],
426            selected: Some(1),
427            sorted: false,
428        };
429
430        unsafe {
431            array.modify_as_c_llama_token_data_array(|c_array| {
432                assert_eq!(c_array.selected, 1);
433            });
434        }
435
436        assert_eq!(array.selected, Some(1));
437    }
438}
llama_cpp_bindings/token/data_array.rs

llama_cpp_bindings/token/
data_array.rs
