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

ferrum_types/
sampling.rs
