Skip to main content

llama_cpp_4/
token.rs

1//! Safe wrappers around `llama_token_data` and `llama_token_data_array`.
2
3use std::fmt::Debug;
4use std::fmt::Display;
5
6pub mod data;
7pub mod data_array;
8
9/// A safe wrapper for `llama_token`.
10///
11/// This struct wraps around a `llama_token` and implements various traits for safe usage, including
12/// `Clone`, `Copy`, `Debug`, `Eq`, `PartialEq`, `Ord`, `PartialOrd`, and `Hash`. The `Display` trait
13/// is also implemented to provide a simple way to format `LlamaToken` for printing.
14#[repr(transparent)]
15#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
16#[allow(clippy::module_name_repetitions)]
17pub struct LlamaToken(pub llama_cpp_sys_4::llama_token);
18
19impl Display for LlamaToken {
20    /// Formats the `LlamaToken` for display by printing its inner value.
21    ///
22    /// This implementation allows you to easily print a `LlamaToken` by using `{}` in formatting macros.
23    ///
24    /// # Example
25    ///
26    /// ```
27    /// # use llama_cpp_4::token::LlamaToken;
28    /// let token = LlamaToken::new(42);
29    /// println!("{}", token); // Prints: 42
30    /// ```
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "{}", self.0)
33    }
34}
35
36impl LlamaToken {
37    /// Creates a new `LlamaToken` from an `i32`.
38    ///
39    /// This constructor allows you to easily create a `LlamaToken` from a raw integer value representing
40    /// the token's ID. This is useful when interacting with external systems that provide token IDs as integers.
41    ///
42    /// # Example
43    ///
44    /// ```
45    /// # use llama_cpp_4::token::LlamaToken;
46    /// let token = LlamaToken::new(0);
47    /// assert_eq!(token, LlamaToken(0));
48    /// ```
49    ///
50    /// # Parameters
51    ///
52    /// - `token_id`: The integer ID for the token.
53    ///
54    /// # Returns
55    ///
56    /// Returns a new instance of `LlamaToken` wrapping the provided `token_id`.
57    #[must_use]
58    pub fn new(token_id: i32) -> Self {
59        Self(token_id)
60    }
61}
62
63/// Converts a vector of `llama_token` to a vector of `LlamaToken` without memory allocation,
64/// and consumes the original vector. This conversion is safe because `LlamaToken` is repr(transparent),
65/// meaning it is just a wrapper around the raw `llama_token` type.
66///
67/// # Safety
68///
69/// This operation is safe because `LlamaToken` has a `repr(transparent)` attribute, ensuring that
70/// the memory layout of `LlamaToken` is the same as that of the underlying `llama_token` type.
71#[must_use]
72pub fn from_vec_token_sys(mut vec_sys: Vec<llama_cpp_sys_4::llama_token>) -> Vec<LlamaToken> {
73    let ptr = vec_sys.as_mut_ptr().cast::<LlamaToken>();
74    unsafe { Vec::from_raw_parts(ptr, vec_sys.len(), vec_sys.capacity()) }
75}
76
77/// Converts a vector of `LlamaToken` to a vector of `llama_token` without memory allocation,
78/// and consumes the original vector. This conversion is safe because `LlamaToken` is repr(transparent),
79/// meaning it is just a wrapper around the raw `llama_token` type.
80///
81/// # Safety
82///
83/// This operation is safe because `LlamaToken` has a `repr(transparent)` attribute, ensuring that
84/// the memory layout of `LlamaToken` is the same as that of the underlying `llama_token` type.
85#[must_use]
86pub fn to_vec_token_sys(mut vec_llama: Vec<LlamaToken>) -> Vec<llama_cpp_sys_4::llama_token> {
87    let ptr = vec_llama
88        .as_mut_ptr()
89        .cast::<llama_cpp_sys_4::llama_token>();
90    unsafe { Vec::from_raw_parts(ptr, vec_llama.len(), vec_llama.capacity()) }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use std::time::Instant;
97
98    #[test]
99    fn test_new_llama_token() {
100        let token = LlamaToken::new(42);
101        assert_eq!(token, LlamaToken(42)); // Verify that the created token has the expected value
102    }
103
104    #[test]
105    fn test_llama_token_display() {
106        let token = LlamaToken::new(99);
107        assert_eq!(format!("{}", token), "99"); // Verify that the token formats correctly
108    }
109
110    #[test]
111    fn test_from_vec_token_sys() {
112        // Test converting a vector of raw `llama_token` to a vector of `LlamaToken`
113        let vec_sys: Vec<llama_cpp_sys_4::llama_token> = vec![1, 2, 3];
114        let vec_llama = from_vec_token_sys(vec_sys);
115
116        // Ensure that the conversion works correctly
117        assert_eq!(vec_llama.len(), 3);
118        assert_eq!(vec_llama[0], LlamaToken(1));
119        assert_eq!(vec_llama[1], LlamaToken(2));
120        assert_eq!(vec_llama[2], LlamaToken(3));
121    }
122
123    #[test]
124    fn test_to_vec_token_sys() {
125        // Test converting a vector of `LlamaToken` to a vector of raw `llama_token`
126        let vec_llama = vec![LlamaToken(10), LlamaToken(20), LlamaToken(30)];
127        let vec_sys = to_vec_token_sys(vec_llama);
128
129        // Ensure that the conversion works correctly
130        assert_eq!(vec_sys.len(), 3);
131        assert_eq!(vec_sys[0], 10);
132        assert_eq!(vec_sys[1], 20);
133        assert_eq!(vec_sys[2], 30);
134    }
135
136    #[test]
137    fn benchmark_to_vec_token_sys() {
138        // Benchmark the speed of to_vec_token_sys by timing it
139        let vec_llama: Vec<LlamaToken> = (0..100_000).map(LlamaToken::new).collect();
140
141        let start = Instant::now();
142        let _vec_sys = to_vec_token_sys(vec_llama);
143        let duration = start.elapsed();
144
145        println!(
146            "Time taken to convert Vec<LlamaToken> to Vec<llama_token>: {:?}",
147            duration
148        );
149
150        // Here we can assert that the conversion took a reasonable amount of time.
151        // This threshold is arbitrary and can be adjusted according to expected performance.
152        assert!(duration.as_micros() < 1_000); // Ensure it takes less than 1ms
153    }
154}