Merge pull request #665 from nhaghighat/add_logit_bias_sampler

Add Logit Sampler

Author: MarcusDunn

llama-cpp-2/src/sampling.rs

@@ -7,6 +7,7 @@ use std::fmt::{Debug, Formatter};
 use crate::context::LlamaContext;
 use crate::model::LlamaModel;
 use crate::token::data_array::LlamaTokenDataArray;
+use crate::token::logit_bias::LlamaLogitBias;
 use crate::token::LlamaToken;
 
 /// A safe wrapper around `llama_sampler`.
@@ -376,6 +377,42 @@ impl LlamaSampler {
         let sampler = unsafe { llama_cpp_sys_2::llama_sampler_init_greedy() };
         Self { sampler }
     }
+
+    /// Creates a sampler that applies bias values to specific tokens during sampling.
+    ///
+    /// # Parameters
+    /// - ``n_vocab``: [`LlamaModel::n_vocab`]
+    /// - ``biases``: Slice of [`LlamaLogitBias`] values specifying token-bias pairs
+    ///
+    /// # Example
+    /// ```rust
+    /// use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
+    /// use llama_cpp_2::sampling::LlamaSampler;
+    ///
+    /// let biases = vec![
+    ///     LlamaLogitBias::new(LlamaToken(1), 1.5),  // Increase probability of token 1
+    ///     LlamaLogitBias::new(LlamaToken(2), -1.0), // Decrease probability of token 2
+    /// ];
+    ///
+    /// // Assuming vocab_size of 32000
+    /// let sampler = LlamaSampler::logit_bias(32000, &biases);
+    /// ```
+    #[must_use]
+    pub fn logit_bias(n_vocab: i32, biases: &[LlamaLogitBias]) -> Self {
+
+        let data = biases.as_ptr().cast::<llama_cpp_sys_2::llama_logit_bias>();
+        
+        let sampler = unsafe {
+            llama_cpp_sys_2::llama_sampler_init_logit_bias(
+                n_vocab,
+                biases.len() as i32,
+                data,
+            )
+        };
+        
+        Self { sampler }
+    }
+
 }
 
 impl Drop for LlamaSampler {

llama-cpp-2/src/token.rs

@@ -5,6 +5,7 @@ use std::fmt::Display;
 
 pub mod data;
 pub mod data_array;
+pub mod logit_bias;
 
 /// A safe wrapper for `llama_token`.
 #[repr(transparent)]

llama-cpp-2/src/token/logit_bias.rs

@@ -0,0 +1,93 @@
+//! Safe wrapper around `llama_logit_bias`.
+use crate::token::LlamaToken;
+
+/// A transparent wrapper around `llama_logit_bias`.
+///
+/// Represents a bias to be applied to a specific token during text generation.
+/// The bias modifies the likelihood of the token being selected.
+///
+/// Do not rely on `repr(transparent)` for this type. It should be considered an implementation
+/// detail and may change across minor versions.
+#[derive(Clone, Copy, Debug, PartialEq)]
+#[repr(transparent)]
+#[allow(clippy::module_name_repetitions)]
+pub struct LlamaLogitBias {
+    logit_bias: llama_cpp_sys_2::llama_logit_bias,
+}
+
+impl LlamaLogitBias {
+    /// Creates a new logit bias for a specific token with the given bias value.
+    /// 
+    /// # Examples
+    /// ```
+    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
+    /// let token = LlamaToken::new(1);
+    /// let bias = LlamaLogitBias::new(token, 1.5);
+    /// ```
+    #[must_use]
+    pub fn new(LlamaToken(token): LlamaToken, bias: f32) -> Self {
+        Self {
+            logit_bias: llama_cpp_sys_2::llama_logit_bias {
+                token,
+                bias,
+            },
+        }
+    }
+
+    /// Gets the token this bias applies to.
+    /// 
+    /// # Examples
+    /// ```
+    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
+    /// let token = LlamaToken::new(1);
+    /// let bias = LlamaLogitBias::new(token, 1.5);
+    /// assert_eq!(bias.token(), token);
+    /// ```
+    #[must_use]
+    pub fn token(&self) -> LlamaToken {
+        LlamaToken(self.logit_bias.token)
+    }
+
+    /// Gets the bias value.
+    /// 
+    /// # Examples
+    /// ```
+    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
+    /// let token = LlamaToken::new(1);
+    /// let bias = LlamaLogitBias::new(token, 1.5);
+    /// assert_eq!(bias.bias(), 1.5);
+    /// ```
+    #[must_use]
+    pub fn bias(&self) -> f32 {
+        self.logit_bias.bias
+    }
+
+    /// Sets the token this bias applies to.
+    /// 
+    /// # Examples
+    /// ```
+    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
+    /// let token = LlamaToken::new(1);
+    /// let mut bias = LlamaLogitBias::new(token, 1.5);
+    /// let new_token = LlamaToken::new(2);
+    /// bias.set_token(new_token);
+    /// assert_eq!(bias.token(), new_token);
+    /// ```
+    pub fn set_token(&mut self, token: LlamaToken) {
+        self.logit_bias.token = token.0;
+    }
+
+    /// Sets the bias value.
+    /// 
+    /// # Examples
+    /// ```
+    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
+    /// let token = LlamaToken::new(1);
+    /// let mut bias = LlamaLogitBias::new(token, 1.5);
+    /// bias.set_bias(2.0);
+    /// assert_eq!(bias.bias(), 2.0);
+    /// ```
+    pub fn set_bias(&mut self, bias: f32) {
+        self.logit_bias.bias = bias;
+    }
+}