token_sampling/
chain.rs

1use crate::candidates::{Candidates, Error};
2use crate::rng::Rng;
3use crate::samplers::{
4    Dist, Dry, DynTemperature, MinP, Penalties, Sampler, SamplerCtx, TailFree, TopK, TopP, Typical,
5    Xtc,
6};
7use core::fmt;
8
9/// Configuration for the standard stage order.
10#[derive(Clone, Debug, Default, PartialEq)]
11pub struct ChainConfig {
12    /// Repetition, frequency, and presence penalties.
13    pub penalties: Penalties,
14    /// Sequence-extension penalty.
15    pub dry: Dry,
16    /// Largest-logit truncation.
17    pub top_k: TopK,
18    /// Tail-free truncation.
19    pub tail_free: TailFree,
20    /// Locally typical truncation.
21    pub typical: Typical,
22    /// Cumulative probability truncation.
23    pub top_p: TopP,
24    /// Maximum-relative probability truncation.
25    pub min_p: MinP,
26    /// Top-choice exclusion.
27    pub xtc: Xtc,
28    /// Fixed or entropy-mapped temperature.
29    pub temperature: DynTemperature,
30}
31
32/// An ordered set of sampling stages with reusable candidate storage.
33pub struct SamplerChain {
34    stages: Vec<Box<dyn Sampler>>,
35    candidates: Candidates,
36}
37
38impl SamplerChain {
39    /// Creates an empty chain.
40    pub const fn new() -> Self {
41        Self {
42            stages: Vec::new(),
43            candidates: Candidates::new(),
44        }
45    }
46
47    /// Appends one concrete stage and returns the chain.
48    pub fn push(mut self, sampler: impl Sampler + 'static) -> Self {
49        self.stages.push(Box::new(sampler));
50        self
51    }
52
53    /// Appends a boxed stage.
54    pub fn push_dyn(&mut self, sampler: Box<dyn Sampler>) {
55        self.stages.push(sampler);
56    }
57
58    /// Builds the standard stage order from `config`.
59    pub fn llama_default(config: ChainConfig) -> Self {
60        Self::new()
61            .push(config.penalties)
62            .push(config.dry)
63            .push(config.top_k)
64            .push(config.tail_free)
65            .push(config.typical)
66            .push(config.top_p)
67            .push(config.min_p)
68            .push(config.xtc)
69            .push(config.temperature)
70            .push(Dist)
71    }
72
73    /// Loads logits, runs every stage, and returns the selected token.
74    pub fn sample(
75        &mut self,
76        logits: &mut [f32],
77        history: &[u32],
78        rng: &mut dyn Rng,
79    ) -> Result<u32, Error> {
80        self.candidates.load(logits)?;
81        let mut context = SamplerCtx { rng, history };
82        for stage in &mut self.stages {
83            stage.apply(&mut self.candidates, &mut context)?;
84        }
85        self.candidates.selected().ok_or(Error::NoSelection)
86    }
87
88    /// Returns the candidate set left by the most recent sample call.
89    pub fn candidates(&self) -> &Candidates {
90        &self.candidates
91    }
92
93    /// Clears selection and per-sequence sampler state.
94    pub fn reset(&mut self) {
95        for stage in &mut self.stages {
96            stage.reset();
97        }
98        self.candidates.selected = None;
99    }
100}
101
102impl fmt::Debug for SamplerChain {
103    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
104        formatter
105            .debug_struct("SamplerChain")
106            .field("stage_count", &self.stages.len())
107            .field("candidates", &self.candidates)
108            .finish()
109    }
110}
111
112impl Default for SamplerChain {
113    fn default() -> Self {
114        Self::new()
115    }
116}
token_sampling/chain.rs

token_sampling/
chain.rs
