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}