babichjacob_llama_cpp_2/
model.rs

1//! A safe wrapper around `llama_model`.
2use std::ffi::CStr;
3use std::ffi::CString;
4use std::num::NonZeroU16;
5use std::os::raw::c_int;
6use std::path::Path;
7use std::ptr::NonNull;
8
9use crate::context::params::LlamaContextParams;
10use crate::context::LlamaContext;
11use crate::llama_backend::LlamaBackend;
12use crate::model::params::LlamaModelParams;
13use crate::token::LlamaToken;
14use crate::token_type::{LlamaTokenAttr, LlamaTokenAttrs};
15use crate::{
16    ApplyChatTemplateError, ChatTemplateError, LlamaContextLoadError, LlamaLoraAdapterInitError,
17    LlamaModelLoadError, NewLlamaChatMessageError, StringToTokenError, TokenToStringError,
18};
19
20pub mod params;
21
22/// A safe wrapper around `llama_model`.
23#[derive(Debug)]
24#[repr(transparent)]
25#[allow(clippy::module_name_repetitions)]
26pub struct LlamaModel {
27    pub(crate) model: NonNull<llama_cpp_sys_2::llama_model>,
28}
29
30/// A safe wrapper around `llama_lora_adapter`.
31#[derive(Debug)]
32#[repr(transparent)]
33#[allow(clippy::module_name_repetitions)]
34pub struct LlamaLoraAdapter {
35    pub(crate) lora_adapter: NonNull<llama_cpp_sys_2::llama_lora_adapter>,
36}
37
38/// A Safe wrapper around `llama_chat_message`
39#[derive(Debug, Eq, PartialEq, Clone)]
40pub struct LlamaChatMessage {
41    role: CString,
42    content: CString,
43}
44
45impl LlamaChatMessage {
46    /// Create a new `LlamaChatMessage`
47    pub fn new(role: String, content: String) -> Result<Self, NewLlamaChatMessageError> {
48        Ok(Self {
49            role: CString::new(role)?,
50            content: CString::new(content)?,
51        })
52    }
53}
54
55/// How to determine if we should prepend a bos token to tokens
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum AddBos {
58    /// Add the beginning of stream token to the start of the string.
59    Always,
60    /// Do not add the beginning of stream token to the start of the string.
61    Never,
62}
63
64/// How to determine if we should tokenize special tokens
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum Special {
67    /// Allow tokenizing special and/or control tokens which otherwise are not exposed and treated as plaintext. Does not insert a leading space.
68    Tokenize,
69    /// Treat special and/or control tokens as plaintext.
70    Plaintext,
71}
72
73unsafe impl Send for LlamaModel {}
74
75unsafe impl Sync for LlamaModel {}
76
77impl LlamaModel {
78    /// get the number of tokens the model was trained on
79    ///
80    /// # Panics
81    ///
82    /// If the number of tokens the model was trained on does not fit into an `u32`. This should be impossible on most
83    /// platforms due to llama.cpp returning a `c_int` (i32 on most platforms) which is almost certainly positive.
84    #[must_use]
85    pub fn n_ctx_train(&self) -> u32 {
86        let n_ctx_train = unsafe { llama_cpp_sys_2::llama_n_ctx_train(self.model.as_ptr()) };
87        u32::try_from(n_ctx_train).expect("n_ctx_train fits into an u32")
88    }
89
90    /// Get all tokens in the model.
91    pub fn tokens(
92        &self,
93        special: Special,
94    ) -> impl Iterator<Item = (LlamaToken, Result<String, TokenToStringError>)> + '_ {
95        (0..self.n_vocab())
96            .map(LlamaToken::new)
97            .map(move |llama_token| (llama_token, self.token_to_str(llama_token, special)))
98    }
99
100    /// Get the beginning of stream token.
101    #[must_use]
102    pub fn token_bos(&self) -> LlamaToken {
103        let token = unsafe { llama_cpp_sys_2::llama_token_bos(self.model.as_ptr()) };
104        LlamaToken(token)
105    }
106
107    /// Get the end of stream token.
108    #[must_use]
109    pub fn token_eos(&self) -> LlamaToken {
110        let token = unsafe { llama_cpp_sys_2::llama_token_eos(self.model.as_ptr()) };
111        LlamaToken(token)
112    }
113
114    /// Get the newline token.
115    #[must_use]
116    pub fn token_nl(&self) -> LlamaToken {
117        let token = unsafe { llama_cpp_sys_2::llama_token_nl(self.model.as_ptr()) };
118        LlamaToken(token)
119    }
120
121    /// Check if a token represents the end of generation (end of turn, end of sequence, etc.)
122    #[must_use]
123    pub fn is_eog_token(&self, token: LlamaToken) -> bool {
124        unsafe { llama_cpp_sys_2::llama_token_is_eog(self.model.as_ptr(), token.0) }
125    }
126
127    /// Get the decoder start token token.
128    #[must_use]
129    pub fn decode_start_token(&self) -> LlamaToken {
130        let token =
131            unsafe { llama_cpp_sys_2::llama_model_decoder_start_token(self.model.as_ptr()) };
132        LlamaToken(token)
133    }
134
135    /// Convert single token to a string.
136    ///
137    /// # Errors
138    ///
139    /// See [`TokenToStringError`] for more information.
140    pub fn token_to_str(
141        &self,
142        token: LlamaToken,
143        special: Special,
144    ) -> Result<String, TokenToStringError> {
145        self.token_to_str_with_size(token, 32, special)
146    }
147
148    /// Convert single token to bytes.
149    ///
150    /// # Errors
151    ///
152    /// See [`TokenToStringError`] for more information.
153    pub fn token_to_bytes(
154        &self,
155        token: LlamaToken,
156        special: Special,
157    ) -> Result<Vec<u8>, TokenToStringError> {
158        self.token_to_bytes_with_size(token, 32, special, None)
159    }
160
161    /// Convert a vector of tokens to a single string.
162    ///
163    /// # Errors
164    ///
165    /// See [`TokenToStringError`] for more information.
166    pub fn tokens_to_str(
167        &self,
168        tokens: &[LlamaToken],
169        special: Special,
170    ) -> Result<String, TokenToStringError> {
171        let mut builder = String::with_capacity(tokens.len() * 4);
172        for str in tokens
173            .iter()
174            .copied()
175            .map(|t| self.token_to_str(t, special))
176        {
177            builder += &str?;
178        }
179        Ok(builder)
180    }
181
182    /// Convert a string to a Vector of tokens.
183    ///
184    /// # Errors
185    ///
186    /// - if [`str`] contains a null byte.
187    ///
188    /// # Panics
189    ///
190    /// - if there is more than [`usize::MAX`] [`LlamaToken`]s in [`str`].
191    ///
192    ///
193    /// ```no_run
194    /// use llama_cpp_2::model::LlamaModel;
195    ///
196    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
197    /// use std::path::Path;
198    /// use llama_cpp_2::model::AddBos;
199    /// let backend = llama_cpp_2::llama_backend::LlamaBackend::init()?;
200    /// let model = LlamaModel::load_from_file(&backend, Path::new("path/to/model"), &Default::default())?;
201    /// let tokens = model.str_to_token("Hello, World!", AddBos::Always)?;
202    /// # Ok(())
203    /// # }
204    pub fn str_to_token(
205        &self,
206        str: &str,
207        add_bos: AddBos,
208    ) -> Result<Vec<LlamaToken>, StringToTokenError> {
209        let add_bos = match add_bos {
210            AddBos::Always => true,
211            AddBos::Never => false,
212        };
213
214        let tokens_estimation = std::cmp::max(8, (str.len() / 2) + usize::from(add_bos));
215        let mut buffer = Vec::with_capacity(tokens_estimation);
216
217        let c_string = CString::new(str)?;
218        let buffer_capacity =
219            c_int::try_from(buffer.capacity()).expect("buffer capacity should fit into a c_int");
220
221        let size = unsafe {
222            llama_cpp_sys_2::llama_tokenize(
223                self.model.as_ptr(),
224                c_string.as_ptr(),
225                c_int::try_from(c_string.as_bytes().len())?,
226                buffer.as_mut_ptr(),
227                buffer_capacity,
228                add_bos,
229                true,
230            )
231        };
232
233        // if we fail the first time we can resize the vector to the correct size and try again. This should never fail.
234        // as a result - size is guaranteed to be positive here.
235        let size = if size.is_negative() {
236            buffer.reserve_exact(usize::try_from(-size).expect("usize's are larger "));
237            unsafe {
238                llama_cpp_sys_2::llama_tokenize(
239                    self.model.as_ptr(),
240                    c_string.as_ptr(),
241                    c_int::try_from(c_string.as_bytes().len())?,
242                    buffer.as_mut_ptr(),
243                    -size,
244                    add_bos,
245                    true,
246                )
247            }
248        } else {
249            size
250        };
251
252        let size = usize::try_from(size).expect("size is positive and usize ");
253
254        // Safety: `size` < `capacity` and llama-cpp has initialized elements up to `size`
255        unsafe { buffer.set_len(size) }
256        Ok(buffer.into_iter().map(LlamaToken).collect())
257    }
258
259    /// Get the type of a token.
260    ///
261    /// # Panics
262    ///
263    /// If the token type is not known to this library.
264    #[must_use]
265    pub fn token_attr(&self, LlamaToken(id): LlamaToken) -> LlamaTokenAttrs {
266        let token_type = unsafe { llama_cpp_sys_2::llama_token_get_attr(self.model.as_ptr(), id) };
267        LlamaTokenAttrs::try_from(token_type).expect("token type is valid")
268    }
269
270    /// Convert a token to a string with a specified buffer size.
271    ///
272    /// Generally you should use [`LlamaModel::token_to_str`] instead as 8 bytes is enough for most words and
273    /// the extra bytes do not really matter.
274    ///
275    /// # Errors
276    ///
277    /// - if the token type is unknown
278    /// - the resultant token is larger than `buffer_size`.
279    /// - the string returend by llama-cpp is not valid utf8.
280    ///
281    /// # Panics
282    ///
283    /// - if `buffer_size` does not fit into a [`c_int`].
284    /// - if the returned size from llama-cpp does not fit into a [`usize`]. (this should never happen)
285    pub fn token_to_str_with_size(
286        &self,
287        token: LlamaToken,
288        buffer_size: usize,
289        special: Special,
290    ) -> Result<String, TokenToStringError> {
291        let bytes = self.token_to_bytes_with_size(token, buffer_size, special, None)?;
292        Ok(String::from_utf8(bytes)?)
293    }
294
295    /// Convert a token to bytes with a specified buffer size.
296    ///
297    /// Generally you should use [`LlamaModel::token_to_bytes`] instead as 8 bytes is enough for most words and
298    /// the extra bytes do not really matter.
299    ///
300    /// # Errors
301    ///
302    /// - if the token type is unknown
303    /// - the resultant token is larger than `buffer_size`.
304    ///
305    /// # Panics
306    ///
307    /// - if `buffer_size` does not fit into a [`c_int`].
308    /// - if the returned size from llama-cpp does not fit into a [`usize`]. (this should never happen)
309    pub fn token_to_bytes_with_size(
310        &self,
311        token: LlamaToken,
312        buffer_size: usize,
313        special: Special,
314        lstrip: Option<NonZeroU16>,
315    ) -> Result<Vec<u8>, TokenToStringError> {
316        if token == self.token_nl() {
317            return Ok(String::from("\n").into_bytes());
318        }
319
320        // unsure what to do with this in the face of the 'special' arg + attr changes
321        let attrs = self.token_attr(token);
322        if attrs.contains(LlamaTokenAttr::Control)
323            && (token == self.token_bos() || token == self.token_eos())
324        {
325            return Ok(Vec::new());
326        } else if attrs.is_empty()
327            || attrs
328                .intersects(LlamaTokenAttr::Unknown | LlamaTokenAttr::Byte | LlamaTokenAttr::Unused)
329        {
330            return Ok(Vec::new());
331        }
332
333        let special = match special {
334            Special::Tokenize => true,
335            Special::Plaintext => false,
336        };
337
338        let string = CString::new(vec![b'*'; buffer_size]).expect("no null");
339        let len = string.as_bytes().len();
340        let len = c_int::try_from(len).expect("length fits into c_int");
341        let buf = string.into_raw();
342        let lstrip = lstrip.map_or(0, |it| i32::from(it.get()));
343        let size = unsafe {
344            llama_cpp_sys_2::llama_token_to_piece(
345                self.model.as_ptr(),
346                token.0,
347                buf,
348                len,
349                lstrip,
350                special,
351            )
352        };
353
354        match size {
355            0 => Err(TokenToStringError::UnknownTokenType),
356            i if i.is_negative() => Err(TokenToStringError::InsufficientBufferSpace(i)),
357            size => {
358                let string = unsafe { CString::from_raw(buf) };
359                let mut bytes = string.into_bytes();
360                let len = usize::try_from(size).expect("size is positive and fits into usize");
361                bytes.truncate(len);
362                Ok(bytes)
363            }
364        }
365    }
366    /// The number of tokens the model was trained on.
367    ///
368    /// This returns a `c_int` for maximum compatibility. Most of the time it can be cast to an i32
369    /// without issue.
370    #[must_use]
371    pub fn n_vocab(&self) -> i32 {
372        unsafe { llama_cpp_sys_2::llama_n_vocab(self.model.as_ptr()) }
373    }
374
375    /// The type of vocab the model was trained on.
376    ///
377    /// # Panics
378    ///
379    /// If llama-cpp emits a vocab type that is not known to this library.
380    #[must_use]
381    pub fn vocab_type(&self) -> VocabType {
382        let vocab_type = unsafe { llama_cpp_sys_2::llama_vocab_type(self.model.as_ptr()) };
383        VocabType::try_from(vocab_type).expect("invalid vocab type")
384    }
385
386    /// This returns a `c_int` for maximum compatibility. Most of the time it can be cast to an i32
387    /// without issue.
388    #[must_use]
389    pub fn n_embd(&self) -> c_int {
390        unsafe { llama_cpp_sys_2::llama_n_embd(self.model.as_ptr()) }
391    }
392
393    /// Get chat template from model.
394    ///
395    /// # Errors
396    ///
397    /// * If the model has no chat template
398    /// * If the chat template is not a valid [`CString`].
399    #[allow(clippy::missing_panics_doc)] // we statically know this will not panic as
400    pub fn get_chat_template(&self, buf_size: usize) -> Result<String, ChatTemplateError> {
401        // longest known template is about 1200 bytes from llama.cpp
402        let chat_temp = CString::new(vec![b'*'; buf_size]).expect("no null");
403        let chat_ptr = chat_temp.into_raw();
404        let chat_name = CString::new("tokenizer.chat_template").expect("no null bytes");
405
406        let ret = unsafe {
407            llama_cpp_sys_2::llama_model_meta_val_str(
408                self.model.as_ptr(),
409                chat_name.as_ptr(),
410                chat_ptr,
411                buf_size,
412            )
413        };
414
415        if ret < 0 {
416            return Err(ChatTemplateError::MissingTemplate(ret));
417        }
418
419        let template_c = unsafe { CString::from_raw(chat_ptr) };
420        let template = template_c.to_str()?;
421
422        let ret: usize = ret.try_into().unwrap();
423        if template.len() < ret {
424            return Err(ChatTemplateError::BuffSizeError(ret + 1));
425        }
426
427        Ok(template.to_owned())
428    }
429
430    /// Loads a model from a file.
431    ///
432    /// # Errors
433    ///
434    /// See [`LlamaModelLoadError`] for more information.
435    #[tracing::instrument(skip_all, fields(params))]
436    pub fn load_from_file(
437        _: &LlamaBackend,
438        path: impl AsRef<Path>,
439        params: &LlamaModelParams,
440    ) -> Result<Self, LlamaModelLoadError> {
441        let path = path.as_ref();
442        debug_assert!(Path::new(path).exists(), "{path:?} does not exist");
443        let path = path
444            .to_str()
445            .ok_or(LlamaModelLoadError::PathToStrError(path.to_path_buf()))?;
446
447        let cstr = CString::new(path)?;
448        let llama_model =
449            unsafe { llama_cpp_sys_2::llama_load_model_from_file(cstr.as_ptr(), params.params) };
450
451        let model = NonNull::new(llama_model).ok_or(LlamaModelLoadError::NullResult)?;
452
453        tracing::debug!(?path, "Loaded model");
454        Ok(LlamaModel { model })
455    }
456
457    /// Initializes a lora adapter from a file.
458    ///
459    /// # Errors
460    ///
461    /// See [`LlamaLoraAdapterInitError`] for more information.
462    pub fn lora_adapter_init(
463        &self,
464        path: impl AsRef<Path>,
465    ) -> Result<LlamaLoraAdapter, LlamaLoraAdapterInitError> {
466        let path = path.as_ref();
467        debug_assert!(Path::new(path).exists(), "{path:?} does not exist");
468
469        let path = path
470            .to_str()
471            .ok_or(LlamaLoraAdapterInitError::PathToStrError(
472                path.to_path_buf(),
473            ))?;
474
475        let cstr = CString::new(path)?;
476        let adapter =
477            unsafe { llama_cpp_sys_2::llama_lora_adapter_init(self.model.as_ptr(), cstr.as_ptr()) };
478
479        let adapter = NonNull::new(adapter).ok_or(LlamaLoraAdapterInitError::NullResult)?;
480
481        tracing::debug!(?path, "Initialized lora adapter");
482        Ok(LlamaLoraAdapter {
483            lora_adapter: adapter,
484        })
485    }
486
487    /// Create a new context from this model.
488    ///
489    /// # Errors
490    ///
491    /// There is many ways this can fail. See [`LlamaContextLoadError`] for more information.
492    // we intentionally do not derive Copy on `LlamaContextParams` to allow llama.cpp to change the type to be non-trivially copyable.
493    #[allow(clippy::needless_pass_by_value)]
494    pub fn new_context(
495        &self,
496        _: &LlamaBackend,
497        params: LlamaContextParams,
498    ) -> Result<LlamaContext, LlamaContextLoadError> {
499        let context_params = params.context_params;
500        let context = unsafe {
501            llama_cpp_sys_2::llama_new_context_with_model(self.model.as_ptr(), context_params)
502        };
503        let context = NonNull::new(context).ok_or(LlamaContextLoadError::NullReturn)?;
504
505        Ok(LlamaContext::new(self, context, params.embeddings()))
506    }
507
508    /// Apply the models chat template to some messages.
509    /// See https://github.com/ggerganov/llama.cpp/wiki/Templates-supported-by-llama_chat_apply_template
510    ///
511    /// `tmpl` of None means to use the default template provided by llama.cpp for the model
512    ///
513    /// # Errors
514    /// There are many ways this can fail. See [`ApplyChatTemplateError`] for more information.
515    #[tracing::instrument(skip_all)]
516    pub fn apply_chat_template(
517        &self,
518        tmpl: Option<String>,
519        chat: Vec<LlamaChatMessage>,
520        add_ass: bool,
521    ) -> Result<String, ApplyChatTemplateError> {
522        // Buffer is twice the length of messages per their recommendation
523        let message_length = chat.iter().fold(0, |acc, c| {
524            acc + c.role.to_bytes().len() + c.content.to_bytes().len()
525        });
526        let mut buff = vec![0; message_length * 4];
527
528        // Build our llama_cpp_sys_2 chat messages
529        let chat: Vec<llama_cpp_sys_2::llama_chat_message> = chat
530            .iter()
531            .map(|c| llama_cpp_sys_2::llama_chat_message {
532                role: c.role.as_ptr(),
533                content: c.content.as_ptr(),
534            })
535            .collect();
536
537        // Set the tmpl pointer
538        let tmpl = tmpl.map(CString::new);
539        let tmpl_ptr = match &tmpl {
540            Some(str) => str.as_ref().map_err(Clone::clone)?.as_ptr(),
541            None => std::ptr::null(),
542        };
543
544        let formatted_chat = unsafe {
545            let res = llama_cpp_sys_2::llama_chat_apply_template(
546                self.model.as_ptr(),
547                tmpl_ptr,
548                chat.as_ptr(),
549                chat.len(),
550                add_ass,
551                buff.as_mut_ptr(),
552                buff.len() as i32,
553            );
554            // A buffer twice the size should be sufficient for all models, if this is not the case for a new model, we can increase it
555            // The error message informs the user to contact a maintainer
556            if res > buff.len() as i32 {
557                return Err(ApplyChatTemplateError::BuffSizeError);
558            }
559            Ok::<String, ApplyChatTemplateError>(
560                CStr::from_ptr(buff.as_mut_ptr())
561                    .to_string_lossy()
562                    .to_string(),
563            )
564        }?;
565        Ok(formatted_chat)
566    }
567}
568
569impl Drop for LlamaModel {
570    fn drop(&mut self) {
571        unsafe { llama_cpp_sys_2::llama_free_model(self.model.as_ptr()) }
572    }
573}
574
575/// a rusty equivalent of `llama_vocab_type`
576#[repr(u32)]
577#[derive(Debug, Eq, Copy, Clone, PartialEq)]
578pub enum VocabType {
579    /// Byte Pair Encoding
580    BPE = llama_cpp_sys_2::LLAMA_VOCAB_TYPE_BPE as _,
581    /// Sentence Piece Tokenizer
582    SPM = llama_cpp_sys_2::LLAMA_VOCAB_TYPE_SPM as _,
583}
584
585/// There was an error converting a `llama_vocab_type` to a `VocabType`.
586#[derive(thiserror::Error, Debug, Eq, PartialEq)]
587pub enum LlamaTokenTypeFromIntError {
588    /// The value is not a valid `llama_token_type`. Contains the int value that was invalid.
589    #[error("Unknown Value {0}")]
590    UnknownValue(llama_cpp_sys_2::llama_vocab_type),
591}
592
593impl TryFrom<llama_cpp_sys_2::llama_vocab_type> for VocabType {
594    type Error = LlamaTokenTypeFromIntError;
595
596    fn try_from(value: llama_cpp_sys_2::llama_vocab_type) -> Result<Self, Self::Error> {
597        match value {
598            llama_cpp_sys_2::LLAMA_VOCAB_TYPE_BPE => Ok(VocabType::BPE),
599            llama_cpp_sys_2::LLAMA_VOCAB_TYPE_SPM => Ok(VocabType::SPM),
600            unknown => Err(LlamaTokenTypeFromIntError::UnknownValue(unknown)),
601        }
602    }
603}
babichjacob_llama_cpp_2/model.rs

babichjacob_llama_cpp_2/
model.rs
