llama_cpp_2/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)]
13#[allow(clippy::module_name_repetitions)]
14pub struct LlamaLogitBias {
15    logit_bias: llama_cpp_sys_2::llama_logit_bias,
16}
17
18impl LlamaLogitBias {
19    /// Creates a new logit bias for a specific token with the given bias value.
20    ///
21    /// # Examples
22    /// ```
23    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
24    /// let token = LlamaToken::new(1);
25    /// let bias = LlamaLogitBias::new(token, 1.5);
26    /// ```
27    #[must_use]
28    pub fn new(LlamaToken(token): LlamaToken, bias: f32) -> Self {
29        Self {
30            logit_bias: llama_cpp_sys_2::llama_logit_bias { token, bias },
31        }
32    }
33
34    /// Gets the token this bias applies to.
35    ///
36    /// # Examples
37    /// ```
38    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
39    /// let token = LlamaToken::new(1);
40    /// let bias = LlamaLogitBias::new(token, 1.5);
41    /// assert_eq!(bias.token(), token);
42    /// ```
43    #[must_use]
44    pub fn token(&self) -> LlamaToken {
45        LlamaToken(self.logit_bias.token)
46    }
47
48    /// Gets the bias value.
49    ///
50    /// # Examples
51    /// ```
52    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
53    /// let token = LlamaToken::new(1);
54    /// let bias = LlamaLogitBias::new(token, 1.5);
55    /// assert_eq!(bias.bias(), 1.5);
56    /// ```
57    #[must_use]
58    pub fn bias(&self) -> f32 {
59        self.logit_bias.bias
60    }
61
62    /// Sets the token this bias applies to.
63    ///
64    /// # Examples
65    /// ```
66    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
67    /// let token = LlamaToken::new(1);
68    /// let mut bias = LlamaLogitBias::new(token, 1.5);
69    /// let new_token = LlamaToken::new(2);
70    /// bias.set_token(new_token);
71    /// assert_eq!(bias.token(), new_token);
72    /// ```
73    pub fn set_token(&mut self, token: LlamaToken) {
74        self.logit_bias.token = token.0;
75    }
76
77    /// Sets the bias value.
78    ///
79    /// # Examples
80    /// ```
81    /// # use llama_cpp_2::token::{LlamaToken, logit_bias::LlamaLogitBias};
82    /// let token = LlamaToken::new(1);
83    /// let mut bias = LlamaLogitBias::new(token, 1.5);
84    /// bias.set_bias(2.0);
85    /// assert_eq!(bias.bias(), 2.0);
86    /// ```
87    pub fn set_bias(&mut self, bias: f32) {
88        self.logit_bias.bias = bias;
89    }
90}
llama_cpp_2/token/logit_bias.rs

llama_cpp_2/token/
logit_bias.rs
