token-sampling 0.1.0

Composable token sampling over f32 logits with caller-supplied randomness
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use crate::candidates::{Candidates, Error};
use crate::rng::Rng;
use crate::samplers::{
    Dist, Dry, DynTemperature, MinP, Penalties, Sampler, SamplerCtx, TailFree, TopK, TopP, Typical,
    Xtc,
};
use core::fmt;

/// Configuration for the standard stage order.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ChainConfig {
    /// Repetition, frequency, and presence penalties.
    pub penalties: Penalties,
    /// Sequence-extension penalty.
    pub dry: Dry,
    /// Largest-logit truncation.
    pub top_k: TopK,
    /// Tail-free truncation.
    pub tail_free: TailFree,
    /// Locally typical truncation.
    pub typical: Typical,
    /// Cumulative probability truncation.
    pub top_p: TopP,
    /// Maximum-relative probability truncation.
    pub min_p: MinP,
    /// Top-choice exclusion.
    pub xtc: Xtc,
    /// Fixed or entropy-mapped temperature.
    pub temperature: DynTemperature,
}

/// An ordered set of sampling stages with reusable candidate storage.
pub struct SamplerChain {
    stages: Vec<Box<dyn Sampler>>,
    candidates: Candidates,
}

impl SamplerChain {
    /// Creates an empty chain.
    pub const fn new() -> Self {
        Self {
            stages: Vec::new(),
            candidates: Candidates::new(),
        }
    }

    /// Appends one concrete stage and returns the chain.
    pub fn push(mut self, sampler: impl Sampler + 'static) -> Self {
        self.stages.push(Box::new(sampler));
        self
    }

    /// Appends a boxed stage.
    pub fn push_dyn(&mut self, sampler: Box<dyn Sampler>) {
        self.stages.push(sampler);
    }

    /// Builds the standard stage order from `config`.
    pub fn llama_default(config: ChainConfig) -> Self {
        Self::new()
            .push(config.penalties)
            .push(config.dry)
            .push(config.top_k)
            .push(config.tail_free)
            .push(config.typical)
            .push(config.top_p)
            .push(config.min_p)
            .push(config.xtc)
            .push(config.temperature)
            .push(Dist)
    }

    /// Loads logits, runs every stage, and returns the selected token.
    pub fn sample(
        &mut self,
        logits: &mut [f32],
        history: &[u32],
        rng: &mut dyn Rng,
    ) -> Result<u32, Error> {
        self.candidates.load(logits)?;
        let mut context = SamplerCtx { rng, history };
        for stage in &mut self.stages {
            stage.apply(&mut self.candidates, &mut context)?;
        }
        self.candidates.selected().ok_or(Error::NoSelection)
    }

    /// Returns the candidate set left by the most recent sample call.
    pub fn candidates(&self) -> &Candidates {
        &self.candidates
    }

    /// Clears selection and per-sequence sampler state.
    pub fn reset(&mut self) {
        for stage in &mut self.stages {
            stage.reset();
        }
        self.candidates.selected = None;
    }
}

impl fmt::Debug for SamplerChain {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("SamplerChain")
            .field("stage_count", &self.stages.len())
            .field("candidates", &self.candidates)
            .finish()
    }
}

impl Default for SamplerChain {
    fn default() -> Self {
        Self::new()
    }
}