llama_cpp_2/token/
data_array.rs

1//! an rusty equivalent of `llama_token_data_array`.
2use std::ptr;
3
4use crate::{sampling::LlamaSampler, token::data::LlamaTokenData};
5
6use super::LlamaToken;
7
8/// a safe wrapper around `llama_token_data_array`.
9#[derive(Debug, Clone, PartialEq)]
10#[allow(clippy::module_name_repetitions)]
11pub struct LlamaTokenDataArray {
12    /// the underlying data
13    pub data: Vec<LlamaTokenData>,
14    /// the index of the selected token in ``data``
15    pub selected: Option<usize>,
16    /// is the data sorted?
17    pub sorted: bool,
18}
19
20impl LlamaTokenDataArray {
21    /// Create a new `LlamaTokenDataArray` from a vector and whether or not the data is sorted.
22    ///
23    /// ```
24    /// # use llama_cpp_2::token::data::LlamaTokenData;
25    /// # use llama_cpp_2::token::data_array::LlamaTokenDataArray;
26    /// # use llama_cpp_2::token::LlamaToken;
27    /// let array = LlamaTokenDataArray::new(vec![
28    ///         LlamaTokenData::new(LlamaToken(0), 0.0, 0.0),
29    ///         LlamaTokenData::new(LlamaToken(1), 0.1, 0.1)
30    ///    ], false);
31    /// assert_eq!(array.data.len(), 2);
32    /// assert_eq!(array.sorted, false);
33    /// ```
34    #[must_use]
35    pub fn new(data: Vec<LlamaTokenData>, sorted: bool) -> Self {
36        Self {
37            data,
38            selected: None,
39            sorted,
40        }
41    }
42
43    /// Create a new `LlamaTokenDataArray` from an iterator and whether or not the data is sorted.
44    /// ```
45    /// # use llama_cpp_2::token::data::LlamaTokenData;
46    /// # use llama_cpp_2::token::data_array::LlamaTokenDataArray;
47    /// # use llama_cpp_2::token::LlamaToken;
48    /// let array = LlamaTokenDataArray::from_iter([
49    ///     LlamaTokenData::new(LlamaToken(0), 0.0, 0.0),
50    ///     LlamaTokenData::new(LlamaToken(1), 0.1, 0.1)
51    /// ], false);
52    /// assert_eq!(array.data.len(), 2);
53    /// assert_eq!(array.sorted, false);
54    pub fn from_iter<T>(data: T, sorted: bool) -> LlamaTokenDataArray
55    where
56        T: IntoIterator<Item = LlamaTokenData>,
57    {
58        Self::new(data.into_iter().collect(), sorted)
59    }
60
61    /// Returns the current selected token, if one exists.
62    #[must_use]
63    pub fn selected_token(&self) -> Option<LlamaToken> {
64        self.data.get(self.selected?).map(LlamaTokenData::id)
65    }
66}
67
68impl LlamaTokenDataArray {
69    /// Modify the underlying data as a `llama_token_data_array`. and reconstruct the `LlamaTokenDataArray`.
70    ///
71    /// # Panics
72    ///
73    /// Panics if some of the safety conditions are not met. (we cannot check all of them at
74    /// runtime so breaking them is UB)
75    ///
76    /// SAFETY:
77    /// The returned array formed by the data pointer and the length must entirely consist of
78    /// initialized token data and the length must be less than the capacity of this array's data
79    /// buffer.
80    /// if the data is not sorted, sorted must be false.
81    pub(crate) unsafe fn modify_as_c_llama_token_data_array<T>(
82        &mut self,
83        modify: impl FnOnce(&mut llama_cpp_sys_2::llama_token_data_array) -> T,
84    ) -> T {
85        let size = self.data.len();
86        let data = self
87            .data
88            .as_mut_ptr()
89            .cast::<llama_cpp_sys_2::llama_token_data>();
90
91        let mut c_llama_token_data_array = llama_cpp_sys_2::llama_token_data_array {
92            data,
93            size,
94            selected: self.selected.and_then(|s| s.try_into().ok()).unwrap_or(-1),
95            sorted: self.sorted,
96        };
97
98        let result = modify(&mut c_llama_token_data_array);
99
100        assert!(
101            c_llama_token_data_array.size <= self.data.capacity(),
102            "Size of the returned array exceeds the data buffer's capacity!"
103        );
104        if !ptr::eq(c_llama_token_data_array.data, data) {
105            ptr::copy(
106                c_llama_token_data_array.data,
107                data,
108                c_llama_token_data_array.size,
109            );
110        }
111        self.data.set_len(c_llama_token_data_array.size);
112
113        self.sorted = c_llama_token_data_array.sorted;
114        self.selected = c_llama_token_data_array
115            .selected
116            .try_into()
117            .ok()
118            .filter(|&s| s < self.data.len());
119
120        result
121    }
122
123    /// Modifies the data array by applying a sampler to it
124    pub fn apply_sampler(&mut self, sampler: &mut LlamaSampler) {
125        unsafe {
126            self.modify_as_c_llama_token_data_array(|c_llama_token_data_array| {
127                llama_cpp_sys_2::llama_sampler_apply(
128                    sampler.sampler.as_mut_ptr(),
129                    c_llama_token_data_array,
130                );
131            });
132        }
133    }
134
135    /// Modifies the data array by applying a sampler to it
136    #[must_use]
137    pub fn with_sampler(mut self, sampler: &mut LlamaSampler) -> Self {
138        self.apply_sampler(sampler);
139        self
140    }
141
142    /// Randomly selects a token from the candidates based on their probabilities.
143    ///
144    /// # Panics
145    /// If the internal llama.cpp sampler fails to select a token.
146    pub fn sample_token(&mut self, seed: u32) -> LlamaToken {
147        self.apply_sampler(&mut LlamaSampler::dist(seed));
148        self.selected_token()
149            .expect("Dist sampler failed to select a token!")
150    }
151
152    /// Selects the token with the highest probability.
153    ///
154    /// # Panics
155    /// If the internal llama.cpp sampler fails to select a token.
156    pub fn sample_token_greedy(&mut self) -> LlamaToken {
157        self.apply_sampler(&mut LlamaSampler::greedy());
158        self.selected_token()
159            .expect("Greedy sampler failed to select a token!")
160    }
161}
llama_cpp_2/token/data_array.rs

llama_cpp_2/token/
data_array.rs
