Skip to main content

babichjacob_llama_cpp_2/token/
data_array.rs

1//! an rusty equivalent of `llama_token_data`.
2use crate::token::data::LlamaTokenData;
3
4/// a safe wrapper around `llama_token_data_array`.
5#[derive(Debug, Clone, PartialEq)]
6#[allow(clippy::module_name_repetitions)]
7pub struct LlamaTokenDataArray {
8    /// the underlying data
9    pub data: Vec<LlamaTokenData>,
10    /// is the data sorted?
11    pub sorted: bool,
12}
13
14impl LlamaTokenDataArray {
15    /// Create a new `LlamaTokenDataArray` from a vector and weather or not the data is sorted.
16    ///
17    /// ```
18    /// # use llama_cpp_2::token::data::LlamaTokenData;
19    /// # use llama_cpp_2::token::data_array::LlamaTokenDataArray;
20    /// # use llama_cpp_2::token::LlamaToken;
21    /// let array = LlamaTokenDataArray::new(vec![
22    ///         LlamaTokenData::new(LlamaToken(0), 0.0, 0.0),
23    ///         LlamaTokenData::new(LlamaToken(1), 0.1, 0.1)
24    ///    ], false);
25    /// assert_eq!(array.data.len(), 2);
26    /// assert_eq!(array.sorted, false);
27    /// ```
28    #[must_use]
29    pub fn new(data: Vec<LlamaTokenData>, sorted: bool) -> Self {
30        Self { data, sorted }
31    }
32
33    /// Create a new `LlamaTokenDataArray` from an iterator and weather or not the data is sorted.
34    /// ```
35    /// # use llama_cpp_2::token::data::LlamaTokenData;
36    /// # use llama_cpp_2::token::data_array::LlamaTokenDataArray;
37    /// # use llama_cpp_2::token::LlamaToken;
38    /// let array = LlamaTokenDataArray::from_iter([
39    ///     LlamaTokenData::new(LlamaToken(0), 0.0, 0.0),
40    ///     LlamaTokenData::new(LlamaToken(1), 0.1, 0.1)
41    /// ], false);
42    /// assert_eq!(array.data.len(), 2);
43    /// assert_eq!(array.sorted, false);
44    pub fn from_iter<T>(data: T, sorted: bool) -> LlamaTokenDataArray
45    where
46        T: IntoIterator<Item = LlamaTokenData>,
47    {
48        Self::new(data.into_iter().collect(), sorted)
49    }
50}