Skip to main content

ferrum_types/
sampling.rs

1//! Sampling and generation parameters
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6use crate::{FerrumError, Result, TokenId};
7
8/// Default repetition penalty for product chat entrypoints.
9///
10/// Chat defaults are greedy, and unpenalized greedy decoding can lock into
11/// deterministic token loops on real models. Keep CLI `run` and OpenAI chat
12/// serving on the same default unless an endpoint exposes an explicit override.
13pub const DEFAULT_CHAT_REPETITION_PENALTY: f32 = 1.1;
14
15/// Sampling parameters for generation
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct SamplingParams {
18    /// Maximum number of tokens to generate
19    pub max_tokens: usize,
20    /// Temperature for randomness (0.0 = deterministic, higher = more random)
21    pub temperature: f32,
22    /// Nucleus sampling probability threshold
23    pub top_p: f32,
24    /// Top-k sampling - consider only top k tokens
25    pub top_k: Option<usize>,
26    /// Repetition penalty to reduce repetitive text
27    pub repetition_penalty: f32,
28    /// Presence penalty for token diversity
29    pub presence_penalty: f32,
30    /// Frequency penalty based on token frequency
31    pub frequency_penalty: f32,
32    /// Stop sequences to end generation
33    pub stop_sequences: Vec<String>,
34    /// Random seed for reproducible generation
35    pub seed: Option<u64>,
36    /// Minimum probability threshold for tokens
37    pub min_p: Option<f32>,
38    /// Tail free sampling parameter
39    pub tfs: Option<f32>,
40    /// Typical sampling parameter
41    pub typical_p: Option<f32>,
42    /// Mirostat sampling parameters
43    pub mirostat: Option<MirostatParams>,
44    /// Response format constraint (JSON mode, schema-constrained, etc.)
45    #[serde(default)]
46    pub response_format: ResponseFormat,
47    /// Point at which a structured-output grammar starts constraining model
48    /// output. Thinking templates can open a reasoning block in the prompt;
49    /// in that case the grammar activates only after the typed delimiter.
50    #[serde(default)]
51    pub structured_output_start: StructuredOutputStart,
52    /// Product completion boundary that must be satisfied before model EOS
53    /// tokens may terminate generation.
54    ///
55    /// This is independent of structured-output grammar activation: ordinary
56    /// text and structured responses need a complete reasoning-to-payload
57    /// transition when the rendered prompt opens one. A complete lexical
58    /// response envelope, such as a tool call, is an alternate terminal path.
59    #[serde(default)]
60    pub response_completion_boundary: ResponseCompletionBoundary,
61}
62
63/// Typed activation boundary for constrained decoding.
64#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
65#[serde(tag = "mode", content = "delimiter", rename_all = "snake_case")]
66pub enum StructuredOutputStart {
67    /// Constrain the first generated token.
68    #[default]
69    Immediate,
70    /// Allow reasoning tokens until this exact tokenizer sequence is emitted,
71    /// then constrain every subsequent token.
72    AfterDelimiter(String),
73}
74
75/// A lexical response envelope that can complete an otherwise pending response.
76///
77/// This describes only token boundaries. The product/API layer remains
78/// responsible for validating the enclosed payload after generation.
79#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
80pub struct ResponseCompletionEnvelope {
81    pub open_token_text: String,
82    pub close_token_text: String,
83    pub max_envelopes: usize,
84}
85
86/// Typed model-EOS boundary for product response completion.
87#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
88#[serde(tag = "mode", rename_all = "snake_case")]
89pub enum ResponseCompletionBoundary {
90    /// Model EOS may terminate any generated token.
91    #[default]
92    Immediate,
93    /// Model EOS remains unavailable until the exact delimiter token sequence
94    /// and at least one subsequent non-whitespace payload token are emitted,
95    /// or until a configured alternate envelope is complete.
96    AfterDelimiterAndPayload {
97        delimiter: String,
98        #[serde(default, skip_serializing_if = "Option::is_none")]
99        alternate_envelope: Option<ResponseCompletionEnvelope>,
100    },
101}
102
103/// Response format for structured output. Mirrors OpenAI's
104/// `response_format` API — no proprietary extensions.
105///
106/// - `Text`: no constraint (default)
107/// - `JsonObject`: output must be a valid JSON object (matches OpenAI's
108///   `{"type": "json_object"}`)
109/// - `JsonSchema(schema)`: output must conform to the given JSON Schema
110///   (matches OpenAI's `{"type": "json_schema", "json_schema": {...}}`).
111///   Internally compiled to a regex FSM for per-token hard masking.
112#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
113#[serde(tag = "type", content = "schema")]
114#[derive(Default)]
115pub enum ResponseFormat {
116    /// No constraint — raw text output.
117    #[default]
118    Text,
119    /// Output must be a valid JSON object.
120    JsonObject,
121    /// Output must conform to the given JSON schema (as a JSON string).
122    JsonSchema(String),
123}
124
125impl Default for SamplingParams {
126    fn default() -> Self {
127        Self {
128            max_tokens: 512,
129            temperature: 1.0,
130            top_p: 1.0,
131            top_k: None,
132            repetition_penalty: 1.0,
133            presence_penalty: 0.0,
134            frequency_penalty: 0.0,
135            stop_sequences: vec![],
136            seed: None,
137            min_p: None,
138            tfs: None,
139            typical_p: None,
140            mirostat: None,
141            response_format: ResponseFormat::default(),
142            structured_output_start: StructuredOutputStart::default(),
143            response_completion_boundary: ResponseCompletionBoundary::default(),
144        }
145    }
146}
147
148impl SamplingParams {
149    /// Create greedy sampling parameters (deterministic)
150    pub fn greedy() -> Self {
151        Self {
152            temperature: 0.0,
153            top_p: 1.0,
154            top_k: None,
155            ..Default::default()
156        }
157    }
158
159    /// Create default sampling parameters with temperature
160    pub fn with_temperature(temperature: f32) -> Self {
161        Self {
162            temperature,
163            ..Default::default()
164        }
165    }
166
167    /// Validate sampling parameters
168    pub fn validate(&self) -> Result<()> {
169        if !self.temperature.is_finite() || self.temperature < 0.0 {
170            return Err(FerrumError::invalid_request(
171                "temperature must be finite and non-negative".to_string(),
172            ));
173        }
174        if !self.top_p.is_finite() || self.top_p <= 0.0 || self.top_p > 1.0 {
175            return Err(FerrumError::invalid_request(
176                "top_p must be in range (0, 1]".to_string(),
177            ));
178        }
179        if let Some(top_k) = self.top_k {
180            if top_k == 0 {
181                return Err(FerrumError::invalid_request(
182                    "top_k must be positive".to_string(),
183                ));
184            }
185        }
186        if !self.repetition_penalty.is_finite() || self.repetition_penalty <= 0.0 {
187            return Err(FerrumError::invalid_request(
188                "repetition_penalty must be finite and positive".to_string(),
189            ));
190        }
191        if !self.presence_penalty.is_finite() || !(-2.0..=2.0).contains(&self.presence_penalty) {
192            return Err(FerrumError::invalid_request(
193                "presence_penalty must be in range [-2, 2]".to_string(),
194            ));
195        }
196        if !self.frequency_penalty.is_finite() || !(-2.0..=2.0).contains(&self.frequency_penalty) {
197            return Err(FerrumError::invalid_request(
198                "frequency_penalty must be in range [-2, 2]".to_string(),
199            ));
200        }
201        if let Some(min_p) = self.min_p {
202            if !min_p.is_finite() || min_p <= 0.0 || min_p > 1.0 {
203                return Err(FerrumError::invalid_request(
204                    "min_p must be in range (0, 1]".to_string(),
205                ));
206            }
207        }
208        if let Some(tfs) = self.tfs {
209            if !tfs.is_finite() || tfs <= 0.0 || tfs > 1.0 {
210                return Err(FerrumError::invalid_request(
211                    "tfs must be in range (0, 1]".to_string(),
212                ));
213            }
214        }
215        if let Some(typical_p) = self.typical_p {
216            if !typical_p.is_finite() || typical_p <= 0.0 || typical_p > 1.0 {
217                return Err(FerrumError::invalid_request(
218                    "typical_p must be in range (0, 1]".to_string(),
219                ));
220            }
221        }
222        if let ResponseCompletionBoundary::AfterDelimiterAndPayload {
223            delimiter,
224            alternate_envelope,
225        } = &self.response_completion_boundary
226        {
227            if delimiter.is_empty() {
228                return Err(FerrumError::invalid_request(
229                    "response completion delimiter must not be empty".to_string(),
230                ));
231            }
232            if let Some(envelope) = alternate_envelope {
233                if envelope.open_token_text.is_empty() || envelope.close_token_text.is_empty() {
234                    return Err(FerrumError::invalid_request(
235                        "response completion envelope tokens must not be empty".to_string(),
236                    ));
237                }
238                if envelope.max_envelopes == 0 {
239                    return Err(FerrumError::invalid_request(
240                        "response completion envelope limit must be greater than zero".to_string(),
241                    ));
242                }
243            }
244        }
245        Ok(())
246    }
247}
248
249/// Mirostat sampling parameters
250#[derive(Debug, Clone, Serialize, Deserialize)]
251pub struct MirostatParams {
252    /// Mirostat mode (1 or 2)
253    pub mode: u8,
254    /// Target entropy
255    pub tau: f32,
256    /// Learning rate
257    pub eta: f32,
258}
259
260/// Sampling presets
261#[derive(Debug, Clone, Serialize, Deserialize)]
262pub struct SamplingPresets {
263    pub presets: HashMap<String, SamplingParams>,
264}
265
266impl Default for SamplingPresets {
267    fn default() -> Self {
268        let mut presets = HashMap::new();
269        presets.insert("greedy".to_string(), SamplingParams::greedy());
270        presets.insert(
271            "creative".to_string(),
272            SamplingParams {
273                temperature: 1.2,
274                top_p: 0.9,
275                top_k: Some(50),
276                repetition_penalty: 1.1,
277                ..Default::default()
278            },
279        );
280        presets.insert(
281            "precise".to_string(),
282            SamplingParams {
283                temperature: 0.3,
284                top_p: 0.95,
285                top_k: Some(20),
286                repetition_penalty: 1.05,
287                ..Default::default()
288            },
289        );
290        Self { presets }
291    }
292}
293
294/// Request priority levels
295#[derive(
296    Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, Default,
297)]
298pub enum Priority {
299    Low = 0,
300    #[default]
301    Normal = 1,
302    High = 2,
303    Critical = 3,
304}
305
306/// Reason for completion
307#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
308pub enum FinishReason {
309    /// Hit maximum token limit
310    Length,
311    /// Hit stop sequence
312    Stop,
313    /// Hit end-of-sequence token
314    EOS,
315    /// Request was cancelled
316    Cancelled,
317    /// Error occurred during generation
318    Error,
319    /// Content filter triggered
320    ContentFilter,
321}
322
323/// Special tokens configuration
324#[derive(Debug, Clone, Serialize, Deserialize, Default)]
325pub struct SpecialTokens {
326    /// Beginning of sequence token
327    pub bos_token: Option<TokenId>,
328    /// End of sequence token
329    pub eos_token: Option<TokenId>,
330    /// Unknown token
331    pub unk_token: Option<TokenId>,
332    /// Padding token
333    pub pad_token: Option<TokenId>,
334    /// Separator token
335    pub sep_token: Option<TokenId>,
336    /// Classification token
337    pub cls_token: Option<TokenId>,
338    /// Mask token
339    pub mask_token: Option<TokenId>,
340    /// Additional end-of-sequence tokens. Models such as Llama-3 and GLM
341    /// declare several `eos_token_id`s in `generation_config.json`;
342    /// `eos_token` holds the primary one and the rest land here.
343    #[serde(default)]
344    pub extra_eos_tokens: Vec<TokenId>,
345}