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}