Skip to main content

Crate token_sampling

Crate token_sampling 

Source
Expand description

Token sampling stages for language-model logits.

The crate accepts f32 logits and token history as slices. A chain owns its candidate storage and reuses it across calls.

§Default chain

use token_sampling::{ChainConfig, SamplerChain, SplitMix64};

let mut chain = SamplerChain::llama_default(ChainConfig::default());
let mut rng = SplitMix64::new(7);
let mut logits = [0.2, 1.4, -0.3, 0.7];
let token = chain.sample(&mut logits, &[1, 2], &mut rng)?;
assert!(token < logits.len() as u32);

§Custom chain

use token_sampling::{Dist, SamplerChain, SplitMix64, Temperature, TopK};

let mut chain = SamplerChain::new()
    .push(TopK { k: 8 })
    .push(Temperature { t: 0.7 })
    .push(Dist);
let mut rng = SplitMix64::new(11);
let mut logits = [0.1, 0.9, 0.4];
let token = chain.sample(&mut logits, &[], &mut rng)?;
assert!(token < 3);

§Mirostat chain

use token_sampling::{MirostatV1, SamplerChain, SplitMix64, Temperature};

let mut chain = SamplerChain::new()
    .push(Temperature { t: 0.8 })
    .push(MirostatV1::new(32));
let mut rng = SplitMix64::new(19);
let mut logits = [0.2; 32];
let token = chain.sample(&mut logits, &[], &mut rng)?;
assert!(token < 32);

Structs§

Candidate
One token candidate with its current logit and probability.
Candidates
Reusable token candidates and stage scratch storage.
ChainConfig
Configuration for the standard stage order.
Dist
Runs softmax and selects one candidate by cumulative probability.
Dry
Applies sequence-extension penalties using recent token repetitions.
DynTemperature
Maps distribution entropy to a temperature range.
Greedy
Selects the first maximum logit in candidate order.
MinP
Keeps logits above a fraction of the maximum token probability.
MirostatV1
Stateful Mirostat version 1 sampling.
MirostatV2
Stateful Mirostat version 2 sampling.
Penalties
Applies repetition, frequency, and presence penalties over recent history.
SamplerChain
An ordered set of sampling stages with reusable candidate storage.
SamplerCtx
Randomness and token history supplied to one sampling stage.
SplitMix64
A small seeded generator based on SplitMix64.
TailFree
Removes the low-probability tail using normalized second differences.
TapeRng
An RNG that replays borrowed random words.
Temperature
Divides logits by a fixed temperature, with greedy masking at zero.
TopK
Keeps the k candidates with the largest logits.
TopP
Keeps the smallest descending prefix that reaches a probability mass.
Typical
Keeps candidates closest to the distribution’s information content.
Xtc
Excludes leading candidates whose probabilities exceed a threshold.

Enums§

Error
Errors returned while loading or sampling logits.

Traits§

Rng
Random bits used by sampling stages.
Sampler
One object-safe stage in a sampling chain.

Functions§

unit
Converts one random word to a value in [0, 1) with 53 bits of precision.