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

ferrum_types/
sampling.rs
