Skip to main content

llama_cpp_4/model/params/
kv_overrides.rs

1//! Key-value overrides for a model.
2
3use crate::model::params::LlamaModelParams;
4use std::ffi::{CStr, CString};
5use std::fmt::Debug;
6
7/// An override value for a model parameter.
8#[derive(Debug, Clone, Copy, PartialEq)]
9pub enum ParamOverrideValue {
10    /// A string value
11    Bool(bool),
12    /// A float value
13    Float(f64),
14    /// A integer value
15    Int(i64),
16    /// A string value
17    Str([std::os::raw::c_char; 128]),
18}
19
20impl ParamOverrideValue {
21    pub(crate) fn tag(&self) -> llama_cpp_sys_4::llama_model_kv_override_type {
22        match self {
23            ParamOverrideValue::Bool(_) => llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_BOOL,
24            ParamOverrideValue::Float(_) => llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_FLOAT,
25            ParamOverrideValue::Int(_) => llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_INT,
26            ParamOverrideValue::Str(_) => llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_STR,
27        }
28    }
29
30    pub(crate) fn value(&self) -> llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 {
31        match self {
32            ParamOverrideValue::Bool(value) => {
33                llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 { val_bool: *value }
34            }
35            ParamOverrideValue::Float(value) => {
36                llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 { val_f64: *value }
37            }
38            ParamOverrideValue::Int(value) => {
39                llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 { val_i64: *value }
40            }
41            ParamOverrideValue::Str(c_string) => {
42                llama_cpp_sys_4::llama_model_kv_override__bindgen_ty_1 { val_str: *c_string }
43            }
44        }
45    }
46}
47
48impl From<&llama_cpp_sys_4::llama_model_kv_override> for ParamOverrideValue {
49    fn from(
50        llama_cpp_sys_4::llama_model_kv_override {
51            key: _,
52            tag,
53            __bindgen_anon_1,
54        }: &llama_cpp_sys_4::llama_model_kv_override,
55    ) -> Self {
56        // SAFETY (all four reads below): `llama_model_kv_override` is a tagged
57        // union — `tag` says which arm of `__bindgen_anon_1` is live, and each
58        // match arm reads exactly the field its tag names. Every variant is a
59        // plain scalar or a fixed `c_char` array, so none carries a pointer
60        // that could dangle.
61        match *tag {
62            llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_INT => {
63                // SAFETY: the tag matched above says this is the live union field.
64                ParamOverrideValue::Int(unsafe { __bindgen_anon_1.val_i64 })
65            }
66            llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_FLOAT => {
67                // SAFETY: the tag matched above says this is the live union field.
68                ParamOverrideValue::Float(unsafe { __bindgen_anon_1.val_f64 })
69            }
70            llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_BOOL => {
71                // SAFETY: the tag matched above says this is the live union field.
72                ParamOverrideValue::Bool(unsafe { __bindgen_anon_1.val_bool })
73            }
74            llama_cpp_sys_4::LLAMA_KV_OVERRIDE_TYPE_STR => {
75                // SAFETY: the tag matched above says this is the live union field.
76                ParamOverrideValue::Str(unsafe { __bindgen_anon_1.val_str })
77            }
78            _ => unreachable!("Unknown tag of {tag}"),
79        }
80    }
81}
82
83/// A struct implementing [`IntoIterator`] over the key-value overrides for a model.
84#[derive(Debug)]
85pub struct KvOverrides<'a> {
86    model_params: &'a LlamaModelParams,
87}
88
89impl KvOverrides<'_> {
90    pub(super) fn new(model_params: &LlamaModelParams) -> KvOverrides<'_> {
91        KvOverrides { model_params }
92    }
93}
94
95impl<'a> IntoIterator for KvOverrides<'a> {
96    // I'm fairly certain this could be written returning by reference, but I'm not sure how to do it safely. I do not
97    // expect this to be a performance bottleneck so the copy should be fine. (let me know if it's not fine!)
98    type Item = (CString, ParamOverrideValue);
99    type IntoIter = KvOverrideValueIterator<'a>;
100
101    fn into_iter(self) -> Self::IntoIter {
102        KvOverrideValueIterator {
103            model_params: self.model_params,
104            current: 0,
105        }
106    }
107}
108
109/// An iterator over the key-value overrides for a model.
110#[derive(Debug)]
111pub struct KvOverrideValueIterator<'a> {
112    model_params: &'a LlamaModelParams,
113    current: usize,
114}
115
116impl Iterator for KvOverrideValueIterator<'_> {
117    type Item = (CString, ParamOverrideValue);
118
119    fn next(&mut self) -> Option<Self::Item> {
120        let overrides = self.model_params.params.kv_overrides;
121        if overrides.is_null() {
122            return None;
123        }
124
125        // SAFETY: llama.cpp seems to guarantee that the last element contains an empty key or is valid. We've checked
126        // the prev one in the last iteration, the next one should be valid or 0 (and thus safe to deref)
127        let current = unsafe { *overrides.add(self.current) };
128
129        if current.key[0] == 0 {
130            return None;
131        }
132
133        let value = ParamOverrideValue::from(&current);
134
135        // SAFETY: `key` is a fixed-size `c_char` array inside the override
136        // struct, and the non-empty check above means it holds at least one
137        // byte; llama.cpp always NUL-terminates it. `to_owned` copies out
138        // before the borrow ends.
139        let key = unsafe { CStr::from_ptr(current.key.as_ptr()).to_owned() };
140
141        self.current += 1;
142        Some((key, value))
143    }
144}