token-sampling 0.1.0

//! 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);
//! # Ok::<(), token_sampling::Error>(())
//! ```
//!
//! # 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);
//! # Ok::<(), token_sampling::Error>(())
//! ```
//!
//! # 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);
//! # Ok::<(), token_sampling::Error>(())
//! ```

#![forbid(unsafe_code)]
#![warn(missing_docs)]

mod candidates;
mod chain;
mod rng;
mod samplers;

pub use candidates::{Candidate, Candidates, Error};
pub use chain::{ChainConfig, SamplerChain};
pub use rng::{unit, Rng, SplitMix64, TapeRng};
pub use samplers::{
    Dist, Dry, DynTemperature, Greedy, MinP, MirostatV1, MirostatV2, Penalties, Sampler,
    SamplerCtx, TailFree, Temperature, TopK, TopP, Typical, Xtc,
};

#[cfg(test)]
mod unit_tests;