Merge pull request #120 from utilityai/kv-overrides

override model values

Author: MarcusDunn

llama-cpp-2/src/model/params.rs

@@ -1,12 +1,113 @@
 //! A safe wrapper around `llama_model_params`.
 
-use std::fmt::Debug;
+use crate::model::params::kv_overrides::KvOverrides;
+use std::ffi::{c_char, CStr};
+use std::fmt::{Debug, Formatter};
+use std::pin::Pin;
+use std::ptr::null;
+
+pub mod kv_overrides;
 
 /// A safe wrapper around `llama_model_params`.
+///
+/// [`T`] is the type of the backing storage for the key-value overrides. Generally it can be left to [`()`] which will
+/// make your life with the borrow checker much easier.
 #[allow(clippy::module_name_repetitions)]
-#[derive(Debug)]
 pub struct LlamaModelParams {
     pub(crate) params: llama_cpp_sys_2::llama_model_params,
+    kv_overrides: Vec<llama_cpp_sys_2::llama_model_kv_override>,
+}
+
+impl Debug for LlamaModelParams {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LlamaModelParams")
+            .field("n_gpu_layers", &self.params.n_gpu_layers)
+            .field("main_gpu", &self.params.main_gpu)
+            .field("vocab_only", &self.params.vocab_only)
+            .field("use_mmap", &self.params.use_mmap)
+            .field("use_mlock", &self.params.use_mlock)
+            .field("kv_overrides", &"vec of kv_overrides")
+            .finish()
+    }
+}
+
+impl LlamaModelParams {
+    /// See [`KvOverrides`]
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use llama_cpp_2::model::params::LlamaModelParams;
+    /// let params = Box::pin(LlamaModelParams::default());
+    /// let kv_overrides = params.kv_overrides();
+    /// let count = kv_overrides.into_iter().count();
+    /// assert_eq!(count, 0);
+    /// ```
+    #[must_use]
+    pub fn kv_overrides(&self) -> KvOverrides {
+        KvOverrides::new(self)
+    }
+
+    /// Appends a key-value override to the model parameters. It must be pinned as this creates a self-referential struct.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use std::ffi::{CStr, CString};
+    /// use std::pin::pin;
+    /// # use llama_cpp_2::model::params::LlamaModelParams;
+    /// # use llama_cpp_2::model::params::kv_overrides::ParamOverrideValue;
+    /// let mut params = pin!(LlamaModelParams::default());
+    /// let key = CString::new("key").expect("CString::new failed");
+    /// params.append_kv_override(&key, ParamOverrideValue::Int(50));
+    ///
+    /// let kv_overrides = params.kv_overrides().into_iter().collect::<Vec<_>>();
+    /// assert_eq!(kv_overrides.len(), 1);
+    ///
+    /// let (k, v) = &kv_overrides[0];
+    /// assert_eq!(v, &ParamOverrideValue::Int(50));
+    ///
+    /// assert_eq!(k.to_bytes(), b"key", "expected key to be 'key', was {:?}", k);
+    /// ```
+    #[allow(clippy::missing_panics_doc)] // panics are just to enforce internal invariants, not user errors
+    pub fn append_kv_override(
+        self: &mut Pin<&mut Self>,
+        key: &CStr,
+        value: kv_overrides::ParamOverrideValue,
+    ) {
+        let kv_override = self
+            .kv_overrides
+            .get_mut(0)
+            .expect("kv_overrides did not have a next allocated");
+
+        assert_eq!(kv_override.key[0], 0, "last kv_override was not empty");
+
+        // There should be some way to do this without iterating over everything.
+        for (i, &c) in key.to_bytes_with_nul().iter().enumerate() {
+            kv_override.key[i] = c_char::try_from(c).expect("invalid character in key");
+        }
+
+        kv_override.tag = value.tag();
+        kv_override.__bindgen_anon_1 = value.value();
+
+        // set to null pointer for panic safety (as push may move the vector, invalidating the pointer)
+        self.params.kv_overrides = null();
+
+        // push the next one to ensure we maintain the iterator invariant of ending with a 0
+        self.kv_overrides
+            .push(llama_cpp_sys_2::llama_model_kv_override {
+                key: [0; 128],
+                tag: 0,
+                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
+                    int_value: 0,
+                },
+            });
+
+        // set the pointer to the (potentially) new vector
+        self.params.kv_overrides = self.kv_overrides.as_ptr();
+
+        eprintln!("saved ptr: {:?}", self.params.kv_overrides);
+    }
 }
 
 impl LlamaModelParams {
@@ -90,8 +191,16 @@ impl LlamaModelParams {
 /// ```
 impl Default for LlamaModelParams {
     fn default() -> Self {
+        let default_params = unsafe { llama_cpp_sys_2::llama_model_default_params() };
         LlamaModelParams {
-            params: unsafe { llama_cpp_sys_2::llama_model_default_params() },
+            params: default_params,
+            kv_overrides: vec![llama_cpp_sys_2::llama_model_kv_override {
+                key: [0; 128],
+                tag: 0,
+                __bindgen_anon_1: llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
+                    int_value: 0,
+                },
+            }],
         }
     }
 }

llama-cpp-2/src/model/params/kv_overrides.rs

@@ -0,0 +1,126 @@
+//! Key-value overrides for a model.
+
+use crate::model::params::LlamaModelParams;
+use std::ffi::{CStr, CString};
+use std::fmt::Debug;
+
+/// An override value for a model parameter.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum ParamOverrideValue {
+    /// A string value
+    Bool(bool),
+    /// A float value
+    Float(f64),
+    /// A integer value
+    Int(i64),
+}
+
+impl ParamOverrideValue {
+    pub(crate) fn tag(&self) -> llama_cpp_sys_2::llama_model_kv_override_type {
+        match self {
+            ParamOverrideValue::Bool(_) => llama_cpp_sys_2::LLAMA_KV_OVERRIDE_TYPE_BOOL,
+            ParamOverrideValue::Float(_) => llama_cpp_sys_2::LLAMA_KV_OVERRIDE_TYPE_FLOAT,
+            ParamOverrideValue::Int(_) => llama_cpp_sys_2::LLAMA_KV_OVERRIDE_TYPE_INT,
+        }
+    }
+
+    pub(crate) fn value(&self) -> llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
+        match self {
+            ParamOverrideValue::Bool(value) => {
+                llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 { bool_value: *value }
+            }
+            ParamOverrideValue::Float(value) => {
+                llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 {
+                    float_value: *value,
+                }
+            }
+            ParamOverrideValue::Int(value) => {
+                llama_cpp_sys_2::llama_model_kv_override__bindgen_ty_1 { int_value: *value }
+            }
+        }
+    }
+}
+
+impl From<&llama_cpp_sys_2::llama_model_kv_override> for ParamOverrideValue {
+    fn from(
+        llama_cpp_sys_2::llama_model_kv_override {
+            key: _,
+            tag,
+            __bindgen_anon_1,
+        }: &llama_cpp_sys_2::llama_model_kv_override,
+    ) -> Self {
+        match *tag {
+            llama_cpp_sys_2::LLAMA_KV_OVERRIDE_TYPE_INT => {
+                ParamOverrideValue::Int(unsafe { __bindgen_anon_1.int_value })
+            }
+            llama_cpp_sys_2::LLAMA_KV_OVERRIDE_TYPE_FLOAT => {
+                ParamOverrideValue::Float(unsafe { __bindgen_anon_1.float_value })
+            }
+            llama_cpp_sys_2::LLAMA_KV_OVERRIDE_TYPE_BOOL => {
+                ParamOverrideValue::Bool(unsafe { __bindgen_anon_1.bool_value })
+            }
+            _ => unreachable!("Unknown tag of {tag}"),
+        }
+    }
+}
+
+/// A struct implementing [`IntoIterator`] over the key-value overrides for a model.
+#[derive(Debug)]
+pub struct KvOverrides<'a> {
+    model_params: &'a LlamaModelParams,
+}
+
+impl KvOverrides<'_> {
+    pub(super) fn new(
+        model_params: &LlamaModelParams,
+    ) -> KvOverrides {
+        KvOverrides { model_params }
+    }
+}
+
+impl<'a> IntoIterator for KvOverrides<'a> {
+    // I'm fairly certain this could be written returning by reference, but I'm not sure how to do it safely. I do not
+    // expect this to be a performance bottleneck so the copy should be fine. (let me know if it's not fine!)
+    type Item = (CString, ParamOverrideValue);
+    type IntoIter = KvOverrideValueIterator<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        KvOverrideValueIterator {
+            model_params: self.model_params,
+            current: 0,
+        }
+    }
+}
+
+/// An iterator over the key-value overrides for a model.
+#[derive(Debug)]
+pub struct KvOverrideValueIterator<'a> {
+    model_params: &'a LlamaModelParams,
+    current: usize,
+}
+
+impl<'a> Iterator for KvOverrideValueIterator<'a> {
+    type Item = (CString, ParamOverrideValue);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let overrides = self.model_params.params.kv_overrides;
+        if overrides.is_null() {
+            return None;
+        }
+
+        // SAFETY: llama.cpp seems to guarantee that the last element contains an empty key or is valid. We've checked
+        // the prev one in the last iteration, the next one should be valid or 0 (and thus safe to deref)
+        let current = unsafe { *overrides.add(self.current) };
+
+        if current.key[0] == 0 {
+            return None;
+        }
+
+        let value = ParamOverrideValue::from(&current);
+
+        let key = unsafe { CStr::from_ptr(current.key.as_ptr()).to_owned() };
+
+        self.current += 1;
+        Some((key, value))
+    }
+}

simple/src/main.rs

@@ -6,7 +6,9 @@
     clippy::cast_sign_loss
 )]
 
-use anyhow::{bail, Context, Result};
+use std::collections::BTreeMap;
+use std::ffi::{CStr, CString};
+use anyhow::{anyhow, bail, Context, Result};
 use clap::Parser;
 use hf_hub::api::sync::ApiBuilder;
 use llama_cpp_2::context::params::LlamaContextParams;
@@ -20,7 +22,10 @@ use llama_cpp_2::token::data_array::LlamaTokenDataArray;
 use std::io::Write;
 use std::num::NonZeroU32;
 use std::path::PathBuf;
+use std::pin::pin;
+use std::str::FromStr;
 use std::time::Duration;
+use llama_cpp_2::model::params::kv_overrides::ParamOverrideValue;
 
 #[derive(clap::Parser, Debug, Clone)]
 struct Args {
@@ -33,12 +38,31 @@ struct Args {
     /// set the length of the prompt + output in tokens
     #[arg(long, default_value_t = 32)]
     n_len: i32,
+    /// override some parameters of the model
+    #[arg(short = 'o', value_parser = parse_key_val)]
+    key_value_overrides: Vec<(String, ParamOverrideValue)>,
     /// Disable offloading layers to the gpu
     #[cfg(feature = "cublas")]
     #[clap(long)]
     disable_gpu: bool,
 }
 
+/// Parse a single key-value pair
+fn parse_key_val(s: &str) -> Result<(String, ParamOverrideValue)> {
+    let pos = s
+        .find('=')
+        .ok_or_else(|| anyhow!("invalid KEY=value: no `=` found in `{}`", s))?;
+    let key = s[..pos].parse()?;
+    let value: String = s[pos + 1..].parse()?;
+    let value = i64::from_str(&value).map(ParamOverrideValue::Int)
+        .or_else(|_| f64::from_str(&value).map(ParamOverrideValue::Float))
+        .or_else(|_| bool::from_str(&value).map(ParamOverrideValue::Bool))
+        .map_err(|_| anyhow!("must be one of i64, f64, or bool"))?;
+    
+    Ok((key, value))
+}
+
+
 #[derive(clap::Subcommand, Debug, Clone)]
 enum Model {
     /// Use an already downloaded model
@@ -79,6 +103,7 @@ fn main() -> Result<()> {
         prompt,
         #[cfg(feature = "cublas")]
         disable_gpu,
+        key_value_overrides,
     } = Args::parse();
 
     // init LLM
@@ -95,6 +120,13 @@ fn main() -> Result<()> {
         #[cfg(not(feature = "cublas"))]
         LlamaModelParams::default()
     };
+    
+    let mut model_params = pin!(model_params);
+    
+    for (k, v) in key_value_overrides.iter() {
+        let k = CString::new(k.as_bytes()).with_context(|| format!("invalid key: {}", k))?;
+        model_params.append_kv_override(k.as_c_str(), *v);
+    }
 
     let model_path = model
         .get_or_load()