token_sampling/lib.rs
1//! Token sampling stages for language-model logits.
2//!
3//! The crate accepts `f32` logits and token history as slices. A chain owns its
4//! candidate storage and reuses it across calls.
5//!
6//! # Default chain
7//!
8//! ```
9//! use token_sampling::{ChainConfig, SamplerChain, SplitMix64};
10//!
11//! let mut chain = SamplerChain::llama_default(ChainConfig::default());
12//! let mut rng = SplitMix64::new(7);
13//! let mut logits = [0.2, 1.4, -0.3, 0.7];
14//! let token = chain.sample(&mut logits, &[1, 2], &mut rng)?;
15//! assert!(token < logits.len() as u32);
16//! # Ok::<(), token_sampling::Error>(())
17//! ```
18//!
19//! # Custom chain
20//!
21//! ```
22//! use token_sampling::{Dist, SamplerChain, SplitMix64, Temperature, TopK};
23//!
24//! let mut chain = SamplerChain::new()
25//! .push(TopK { k: 8 })
26//! .push(Temperature { t: 0.7 })
27//! .push(Dist);
28//! let mut rng = SplitMix64::new(11);
29//! let mut logits = [0.1, 0.9, 0.4];
30//! let token = chain.sample(&mut logits, &[], &mut rng)?;
31//! assert!(token < 3);
32//! # Ok::<(), token_sampling::Error>(())
33//! ```
34//!
35//! # Mirostat chain
36//!
37//! ```
38//! use token_sampling::{MirostatV1, SamplerChain, SplitMix64, Temperature};
39//!
40//! let mut chain = SamplerChain::new()
41//! .push(Temperature { t: 0.8 })
42//! .push(MirostatV1::new(32));
43//! let mut rng = SplitMix64::new(19);
44//! let mut logits = [0.2; 32];
45//! let token = chain.sample(&mut logits, &[], &mut rng)?;
46//! assert!(token < 32);
47//! # Ok::<(), token_sampling::Error>(())
48//! ```
49
50#![forbid(unsafe_code)]
51#![warn(missing_docs)]
52
53mod candidates;
54mod chain;
55mod rng;
56mod samplers;
57
58pub use candidates::{Candidate, Candidates, Error};
59pub use chain::{ChainConfig, SamplerChain};
60pub use rng::{unit, Rng, SplitMix64, TapeRng};
61pub use samplers::{
62 Dist, Dry, DynTemperature, Greedy, MinP, MirostatV1, MirostatV2, Penalties, Sampler,
63 SamplerCtx, TailFree, Temperature, TopK, TopP, Typical, Xtc,
64};
65
66#[cfg(test)]
67mod unit_tests;