Skip to main content

llama_cpp_bindings/token/
logit_bias.rs

1//! Safe wrapper around `llama_logit_bias`.
2use crate::token::LlamaToken;
3
4/// A transparent wrapper around `llama_logit_bias`.
5///
6/// Represents a bias to be applied to a specific token during text generation.
7/// The bias modifies the likelihood of the token being selected.
8///
9/// Do not rely on `repr(transparent)` for this type. It should be considered an implementation
10/// detail and may change across minor versions.
11#[derive(Clone, Copy, Debug, PartialEq)]
12#[repr(transparent)]
13pub struct LlamaLogitBias {
14    logit_bias: llama_cpp_bindings_sys::llama_logit_bias,
15}
16
17impl LlamaLogitBias {
18    /// Creates a new logit bias for a specific token with the given bias value.
19    ///
20    /// # Examples
21    /// ```
22    /// # use llama_cpp_bindings::token::{LlamaToken, logit_bias::LlamaLogitBias};
23    /// let token = LlamaToken::new(1);
24    /// let bias = LlamaLogitBias::new(token, 1.5);
25    /// ```
26    #[must_use]
27    pub fn new(LlamaToken(token): LlamaToken, bias: f32) -> Self {
28        Self {
29            logit_bias: llama_cpp_bindings_sys::llama_logit_bias { token, bias },
30        }
31    }
32
33    /// Gets the token this bias applies to.
34    ///
35    /// # Examples
36    /// ```
37    /// # use llama_cpp_bindings::token::{LlamaToken, logit_bias::LlamaLogitBias};
38    /// let token = LlamaToken::new(1);
39    /// let bias = LlamaLogitBias::new(token, 1.5);
40    /// assert_eq!(bias.token(), token);
41    /// ```
42    #[must_use]
43    pub fn token(&self) -> LlamaToken {
44        LlamaToken(self.logit_bias.token)
45    }
46
47    /// Gets the bias value.
48    ///
49    /// # Examples
50    /// ```
51    /// # use llama_cpp_bindings::token::{LlamaToken, logit_bias::LlamaLogitBias};
52    /// let token = LlamaToken::new(1);
53    /// let bias = LlamaLogitBias::new(token, 1.5);
54    /// assert_eq!(bias.bias(), 1.5);
55    /// ```
56    #[must_use]
57    pub fn bias(&self) -> f32 {
58        self.logit_bias.bias
59    }
60
61    /// Sets the token this bias applies to.
62    ///
63    /// # Examples
64    /// ```
65    /// # use llama_cpp_bindings::token::{LlamaToken, logit_bias::LlamaLogitBias};
66    /// let token = LlamaToken::new(1);
67    /// let mut bias = LlamaLogitBias::new(token, 1.5);
68    /// let new_token = LlamaToken::new(2);
69    /// bias.set_token(new_token);
70    /// assert_eq!(bias.token(), new_token);
71    /// ```
72    pub fn set_token(&mut self, token: LlamaToken) {
73        self.logit_bias.token = token.0;
74    }
75
76    /// Sets the bias value.
77    ///
78    /// # Examples
79    /// ```
80    /// # use llama_cpp_bindings::token::{LlamaToken, logit_bias::LlamaLogitBias};
81    /// let token = LlamaToken::new(1);
82    /// let mut bias = LlamaLogitBias::new(token, 1.5);
83    /// bias.set_bias(2.0);
84    /// assert_eq!(bias.bias(), 2.0);
85    /// ```
86    pub fn set_bias(&mut self, bias: f32) {
87        self.logit_bias.bias = bias;
88    }
89}