token-sampling

token-sampling applies sampling stages to a slice of f32 logits. It includes temperature, top-k, top-p, min-p, locally typical, tail-free, XTC, token penalties, DRY, Mirostat, greedy selection, and a final probability draw.

The caller owns token history and randomness. A chain keeps its candidate buffers between calls, so repeated sampling at one vocabulary size reuses allocated memory.

Example

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

let mut sampler = SamplerChain::llama_default(ChainConfig::default());
let mut rng = SplitMix64::new(42);
let mut logits = [0.3, 1.2, -0.8, 0.6];

let token = sampler.sample(&mut logits, &[3, 1, 3], &mut rng)?;
assert!(token < logits.len() as u32);
# Ok::<(), token_sampling::Error>(())

Build a different order with SamplerChain::push:

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

let sampler = SamplerChain::new()
    .push(TopK { k: 20 })
    .push(Temperature { t: 0.7 })
    .push(Dist);

Rng requires one method, next_u64. The crate ships SplitMix64 for seeded sampling and TapeRng for replaying fixed draws. Enable the rand_core feature to use a rand_core::RngCore implementation directly.

Errors

Candidates::load rejects empty logits, the first NaN or positive infinity, and an all-masked slice. Negative infinity is a valid mask. A chain without a terminal stage returns NoSelection. Validated finite logits produce a token even when a configured stage would overflow f32 arithmetic.

Minimum Rust version

The crate supports Rust 1.63 and later.

License

MIT