llm_samplers/
lib.rs

1//! # LLM Samplers
2//!
3//! Rusty samplers for large language models.
4//!
5//! ## Examples
6//!
7//! You probably won't usually want to use individual [crate::types::Sampler]s. The most typical
8//! use case is going to be chaining a number of samplers together. For reference, here's a link to
9//! the [crate::samplers] module.
10//!
11//! A simple example of constructing a [crate::prelude::SamplerChain]:
12//!
13//! ```rust
14//! use anyhow::Result;
15//!
16//! use llm_samplers::prelude::*;
17//!
18//! pub fn test_chain1() -> Result<()> {
19//!
20//!     let mut logits = Logits::try_from_iter([0.1f32, 0.2, 0.3, 0.4].into_iter())?;
21//!
22//!     // Demonstrating the different ways you can build a SamplerChain.
23//!     // These are all equivalent.
24//!     let mut sc = SamplerChain::new()
25//!         + SampleFlatBias::new([(3, f32::NEG_INFINITY)]);
26//!     sc += SampleTemperature::new(0.8);
27//!     sc.push_sampler(SampleGreedy::new());
28//!
29//!     assert_eq!(
30//!         sc.sample_token(
31//!             // These samplers don't actually need any resources.
32//!             &mut NilSamplerResources::default(),
33//!             &mut logits)?,
34//!         Some(1)
35//!     );
36//!
37//!     // () also implements HasSamplerResources<TokenId = u32>
38//!     // so you could use &mut () here.
39//!     assert_eq!(sc.sample_token(&mut (), &mut logits)?, Some(1));
40//!     Ok(())
41//! }
42//! ```
43//!
44//! Now let's look at something a bit more complicated:
45//!
46//! ```rust
47//! # pub mod example {
48//! use std::sync::{Arc, RwLock};
49//!
50//! use anyhow::Result;
51//! use rand::{SeedableRng, rngs::StdRng};
52//!
53//! use llm_samplers::prelude::*;
54//!
55//! fn test_chain2() -> Result<()> {
56//!     let example_logits = vec![0.1f32, 0.2, 0.3, 0.4];
57//!     let mut res = SimpleSamplerResources::new(
58//!         // Optionally include an RNG resource.
59//!         Some(Box::new(StdRng::seed_from_u64(123))),
60//!         // Optionally include a last tokens resource.
61//!         Some(vec![]),
62//!     );
63//!     let mut logits = Logits::try_from_iter(example_logits.into_iter())?;
64//!     let mut logits2 = logits.clone();
65//!
66//!     // SamplerChain with u32 token id type and f32 logit type.
67//!     let mut sc = SamplerChain::new()
68//!         // Bias logits (this example sets bias for token id 3 to -inf)
69//!         + SampleFlatBias::new([(3, f32::NEG_INFINITY)])
70//!         // Apply a repetition penalty.
71//!         + SampleRepetition::new(1.1, 64)
72//!         // Apply frequency and presence penalties.
73//!         + SampleFreqPresence::new(0.05, 0.1, 64)
74//!         // Apply temperature to logits.
75//!         + SampleTemperature::new(0.8)
76//!         // Sample a token using Mirostat1
77//!         + SampleMirostat1::new(4, 5.0, 0.1);
78//!
79//!     // Put a value into `last_tokens`, this simulates us having already picked
80//!     // that token (3) previously.
81//!     res.with_last_tokens_mut(&mut |tokens| tokens.push(3u32))?;
82//!
83//!     assert_eq!(sc.sample_token(&mut res, &mut logits)?, Some(2));
84//!
85//!     // Now add the last selected token to the list.
86//!     res.with_last_tokens_mut(&mut |tokens| tokens.push(2u32))?;
87//!
88//!     // And pick the next one. *Important*: Note that we don't reuse `logits`.
89//!     // This is because `logits` already has all the filtering/sorting/permutation
90//!     // from the previous sample call applied to it.
91//!     assert_eq!(sc.sample_token(&mut res, &mut logits2)?, Some(1));
92//!     Ok(())
93//! }
94//!
95//! # }
96//! ```
97//!
98//! ## Suggested chains/ordering
99//! Suggestions based on the way `llama.cpp` does it.
100//!
101//! Note that you may not get meaningful results if you do weird stuff like chaining
102//! multiple token selecting samplers. Based on available information, it also seems
103//! like combining Mirostat samplers with top-K, top-P, etc will not work.
104//!
105//! ### Temperature sampling
106//! 1. [SampleFlatBias](crate::samplers::SampleFlatBias) (optional)
107//! 2. [SampleRepetition](crate::samplers::SampleRepetition) (optional)
108//! 3. [SampleFreqPresence](crate::samplers::SampleFreqPresence) (optional)
109//! 4. [SampleTopK](crate::samplers::SampleTopK) (optional)
110//! 5. [SampleTailFree](crate::samplers::SampleTailFree) (optional)
111//! 6. [SampleLocallyTypical](crate::samplers::SampleLocallyTypical) (optional)
112//! 7. [SampleTopP](crate::samplers::SampleTopP) (optional)
113//! 8. [SampleMinP](crate::samplers::SampleMinP) (optional)
114//! 9. [SampleTemperature](crate::samplers::SampleTemperature) (optional)
115//! 10. [SampleRandDistrib](crate::samplers::SampleRandDistrib)
116//!
117//! ### Mirostat V1/V2
118//! 1. [SampleFlatBias](crate::samplers::SampleFlatBias) (optional)
119//! 2. [SampleRepetition](crate::samplers::SampleRepetition) (optional)
120//! 3. [SampleFreqPresence](crate::samplers::SampleFreqPresence) (optional)
121//! 4. [SampleTemperature](crate::samplers::SampleTemperature) (optional)
122//! 5. [SampleMirostat1](crate::samplers::SampleMirostat1) **or** [SampleMirostat2](crate::samplers::SampleMirostat2)
123
124/// # Samplers live here!
125///
126/// Visiting the structs below is probably going to be more helpful than going to the modules.
127pub mod samplers;
128
129/// Types and traits.
130pub mod types;
131
132/// Sampler chains
133mod chain;
134
135/// Sampler resources
136mod resource;
137
138/// Configuring sampler options
139pub mod configure;
140
141#[cfg(test)]
142mod tests;
143
144/// Convenient rexports. The simplest way to use the crate is to just throw a
145/// `use llm_samplers::prelude::*;`
146/// into your project.
147pub mod prelude {
148    #[doc(inline)]
149    pub use crate::{
150        chain::*,
151        configure::{ConfigurableSampler, ConfigureSamplerError},
152        resource::*,
153        samplers::*,
154        types::*,
155    };
156}
llm_samplers/lib.rs

llm_samplers/
lib.rs
