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}