llama_cpp_4/model/
params.rs

1//! A safe wrapper around `llama_model_params`.
2
3use crate::model::params::kv_overrides::KvOverrides;
4use std::ffi::{c_char, CStr};
5use std::fmt::{Debug, Formatter};
6use std::pin::Pin;
7use std::ptr::null;
8
9pub mod kv_overrides;
10
11/// Exact model-file loading strategy exposed by llama.cpp.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13#[repr(u32)]
14pub enum LlamaLoadMode {
15    /// No memory mapping, locking, or direct I/O.
16    None = llama_cpp_sys_4::LLAMA_LOAD_MODE_NONE,
17    /// Memory-map model files when supported.
18    Mmap = llama_cpp_sys_4::LLAMA_LOAD_MODE_MMAP,
19    /// Read model files normally and lock loaded pages in memory.
20    Mlock = llama_cpp_sys_4::LLAMA_LOAD_MODE_MLOCK,
21    /// Memory-map model files and lock mapped pages in memory.
22    MmapMlock = llama_cpp_sys_4::LLAMA_LOAD_MODE_MMAP_MLOCK,
23    /// Use direct I/O when supported.
24    DirectIo = llama_cpp_sys_4::LLAMA_LOAD_MODE_DIRECT_IO,
25}
26
27/// A safe wrapper around `llama_model_params`.
28#[allow(clippy::module_name_repetitions)]
29pub struct LlamaModelParams {
30    pub(crate) params: llama_cpp_sys_4::llama_model_params,
31    kv_overrides: Vec<llama_cpp_sys_4::llama_model_kv_override>,
32}
33
34impl Debug for LlamaModelParams {
35    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
36        f.debug_struct("LlamaModelParams")
37            .field("n_gpu_layers", &self.params.n_gpu_layers)
38            .field("main_gpu", &self.params.main_gpu)
39            .field("vocab_only", &self.params.vocab_only)
40            .field("load_mode", &self.load_mode())
41            .field("kv_overrides", &"vec of kv_overrides")
42            .finish()
43    }
44}
45
46impl LlamaModelParams {
47    /// See [`KvOverrides`]
48    ///
49    /// # Examples
50    ///
51    /// ```rust
52    /// # use llama_cpp_4::model::params::LlamaModelParams;
53    /// let params = Box::pin(LlamaModelParams::default());
54    /// let kv_overrides = params.kv_overrides();
55    /// let count = kv_overrides.into_iter().count();
56    /// assert_eq!(count, 0);
57    /// ```
58    #[must_use]
59    pub fn kv_overrides(&self) -> KvOverrides<'_> {
60        KvOverrides::new(self)
61    }
62
63    /// Appends a key-value override to the model parameters. It must be pinned as this creates a self-referential struct.
64    ///
65    /// # Examples
66    ///
67    /// ```rust
68    /// # use std::ffi::{CStr, CString};
69    /// use std::pin::pin;
70    /// # use llama_cpp_4::model::params::LlamaModelParams;
71    /// # use llama_cpp_4::model::params::kv_overrides::ParamOverrideValue;
72    /// let mut params = pin!(LlamaModelParams::default());
73    /// let key = CString::new("key").expect("CString::new failed");
74    /// params.as_mut().append_kv_override(&key, ParamOverrideValue::Int(50));
75    ///
76    /// let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
77    /// assert_eq!(kv_overrides.len(), 1);
78    ///
79    /// let (k, v) = &kv_overrides[0];
80    /// assert_eq!(v, &ParamOverrideValue::Int(50));
81    ///
82    /// assert_eq!(k.to_bytes(), b"key", "expected key to be 'key', was {:?}", k);
83    /// ```
84    #[allow(clippy::missing_panics_doc)] // panics are just to enforce internal invariants, not user errors
85    pub fn append_kv_override(
86        mut self: Pin<&mut Self>,
87        key: &CStr,
88        value: kv_overrides::ParamOverrideValue,
89    ) {
90        let kv_override = self
91            .kv_overrides
92            .get_mut(0)
93            .expect("kv_overrides did not have a next allocated");
94
95        assert_eq!(kv_override.key[0], 0, "last kv_override was not empty");
96
97        // There should be some way to do this without iterating over everything.
98        for (i, &c) in key.to_bytes_with_nul().iter().enumerate() {
99            kv_override.key[i] = c_char::try_from(c).expect("invalid character in key");
100        }
101
102        kv_override.tag = value.tag();
103        kv_override.__bindgen_anon_1 = value.value();
104
105        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
106        self.params.kv_overrides = null();
107
108        // push the next one to ensure we maintain the iterator invariant of ending with a 0
109        self.kv_overrides
110            .push(llama_cpp_sys_4::llama_model_kv_override {
111                key: [0; 128],
112                tag: 0,
113                __bindgen_anon_1: llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 {
114                    val_i64: 0,
115                },
116            });
117
118        // set the pointer to the (potentially) new vector
119        self.params.kv_overrides = self.kv_overrides.as_ptr();
120
121        eprintln!("saved ptr: {:?}", self.params.kv_overrides);
122    }
123}
124
125impl LlamaModelParams {
126    /// Get the number of layers to offload to the GPU.
127    #[must_use]
128    pub fn n_gpu_layers(&self) -> i32 {
129        self.params.n_gpu_layers
130    }
131
132    /// The GPU that is used for scratch and small tensors
133    #[must_use]
134    pub fn main_gpu(&self) -> i32 {
135        self.params.main_gpu
136    }
137
138    /// only load the vocabulary, no weights
139    #[must_use]
140    pub fn vocab_only(&self) -> bool {
141        self.params.vocab_only
142    }
143
144    /// Returns the exact model-file loading strategy.
145    #[must_use]
146    pub fn load_mode(&self) -> LlamaLoadMode {
147        match self.params.load_mode {
148            llama_cpp_sys_4::LLAMA_LOAD_MODE_MMAP => LlamaLoadMode::Mmap,
149            llama_cpp_sys_4::LLAMA_LOAD_MODE_MLOCK => LlamaLoadMode::Mlock,
150            llama_cpp_sys_4::LLAMA_LOAD_MODE_MMAP_MLOCK => LlamaLoadMode::MmapMlock,
151            llama_cpp_sys_4::LLAMA_LOAD_MODE_DIRECT_IO => LlamaLoadMode::DirectIo,
152            _ => LlamaLoadMode::None,
153        }
154    }
155
156    /// use mmap if possible
157    #[must_use]
158    pub fn use_mmap(&self) -> bool {
159        matches!(
160            self.load_mode(),
161            LlamaLoadMode::Mmap | LlamaLoadMode::MmapMlock
162        )
163    }
164
165    /// force system to keep model in RAM
166    #[must_use]
167    pub fn use_mlock(&self) -> bool {
168        matches!(
169            self.load_mode(),
170            LlamaLoadMode::Mlock | LlamaLoadMode::MmapMlock
171        )
172    }
173
174    /// sets the number of gpu layers to offload to the GPU.
175    /// ```
176    /// # use llama_cpp_4::model::params::LlamaModelParams;
177    /// let params = LlamaModelParams::default();
178    /// let params = params.with_n_gpu_layers(1);
179    /// assert_eq!(params.n_gpu_layers(), 1);
180    /// ```
181    #[must_use]
182    pub fn with_n_gpu_layers(mut self, n_gpu_layers: u32) -> Self {
183        // The only way this conversion can fail is if u32 overflows the i32 - in which case we set
184        // to MAX
185        let n_gpu_layers = i32::try_from(n_gpu_layers).unwrap_or(i32::MAX);
186        self.params.n_gpu_layers = n_gpu_layers;
187        self
188    }
189
190    /// sets the main GPU
191    #[must_use]
192    pub fn with_main_gpu(mut self, main_gpu: i32) -> Self {
193        self.params.main_gpu = main_gpu;
194        self
195    }
196
197    /// sets `vocab_only`
198    #[must_use]
199    pub fn with_vocab_only(mut self, vocab_only: bool) -> Self {
200        self.params.vocab_only = vocab_only;
201        self
202    }
203
204    /// Sets the exact model-file loading strategy.
205    #[must_use]
206    pub fn with_load_mode(mut self, load_mode: LlamaLoadMode) -> Self {
207        self.params.load_mode = load_mode as llama_cpp_sys_4::llama_load_mode;
208        self
209    }
210
211    /// sets `use_mlock`
212    #[must_use]
213    pub fn with_use_mlock(mut self, use_mlock: bool) -> Self {
214        let load_mode = match (self.use_mmap(), use_mlock) {
215            (true, true) => LlamaLoadMode::MmapMlock,
216            (true, false) => LlamaLoadMode::Mmap,
217            (false, true) => LlamaLoadMode::Mlock,
218            (false, false) => LlamaLoadMode::None,
219        };
220        self.params.load_mode = load_mode as llama_cpp_sys_4::llama_load_mode;
221        self
222    }
223}
224
225/// Default parameters for `LlamaModel`. (as defined in llama.cpp by `llama_model_default_params`)
226/// ```
227/// # use llama_cpp_4::model::params::LlamaModelParams;
228/// let params = LlamaModelParams::default();
229/// assert_eq!(params.n_gpu_layers(), -1, "n_gpu_layers should be -1 (all layers)");
230/// assert_eq!(params.main_gpu(), 0, "main_gpu should be 0");
231/// assert_eq!(params.vocab_only(), false, "vocab_only should be false");
232/// assert_eq!(params.use_mmap(), true, "use_mmap should be true");
233/// assert_eq!(params.use_mlock(), false, "use_mlock should be false");
234/// ```
235impl Default for LlamaModelParams {
236    fn default() -> Self {
237        let default_params = unsafe { llama_cpp_sys_4::llama_model_default_params() };
238        LlamaModelParams {
239            params: default_params,
240            // push the next one to ensure we maintain the iterator invariant of ending with a 0
241            kv_overrides: vec![llama_cpp_sys_4::llama_model_kv_override {
242                key: [0; 128],
243                tag: 0,
244                __bindgen_anon_1: llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 {
245                    val_i64: 0,
246                },
247            }],
248        }
249    }
250}
llama_cpp_4/model/params.rs

llama_cpp_4/model/
params.rs
