infotheory/api/
generation.rs

1//! Generation-focused public API surface.
2
3use super::context::RateBackendSession;
4use super::types::{GenerationConfig, GenerationStrategy};
5use crate::error::{InfotheoryError, InfotheoryResult};
6use crate::spec::CompiledRateBackend;
7
8use crate::with_default_ctx;
9
10pub(crate) struct GenerationRng {
11    state: u64,
12}
13
14impl GenerationRng {
15    pub(crate) fn new(seed: u64) -> Self {
16        Self {
17            state: if seed == 0 {
18                0xD00D_F00D_CAFE_BABEu64
19            } else {
20                seed
21            },
22        }
23    }
24
25    fn next_u64(&mut self) -> u64 {
26        let mut x = self.state;
27        x ^= x << 13;
28        x ^= x >> 7;
29        x ^= x << 17;
30        self.state = x;
31        x
32    }
33
34    fn next_f64(&mut self) -> f64 {
35        unit_interval_from_u64(self.next_u64())
36    }
37}
38
39#[inline(always)]
40fn unit_interval_from_u64(bits: u64) -> f64 {
41    ((bits >> 11) as f64) * (1.0 / ((1u64 << 53) as f64))
42}
43
44#[inline(always)]
45fn argmax_log_prob_byte(logps: &[f64; 256]) -> u8 {
46    let mut best_idx = 0usize;
47    let mut best = f64::NEG_INFINITY;
48    for (idx, &logp) in logps.iter().enumerate() {
49        let score = if logp.is_finite() {
50            logp
51        } else {
52            f64::NEG_INFINITY
53        };
54        if score > best {
55            best = score;
56            best_idx = idx;
57        }
58    }
59    best_idx as u8
60}
61
62pub(crate) fn pick_generated_byte(
63    logps: &[f64; 256],
64    config: GenerationConfig,
65    rng: &mut GenerationRng,
66) -> u8 {
67    if matches!(config.strategy, GenerationStrategy::Greedy)
68        || !config.temperature.is_finite()
69        || config.temperature <= 0.0
70    {
71        return argmax_log_prob_byte(logps);
72    }
73
74    let mut entries = [(0u8, f64::NEG_INFINITY); 256];
75    for (idx, &logp) in logps.iter().enumerate() {
76        let scaled = if logp.is_finite() {
77            logp / config.temperature
78        } else {
79            f64::NEG_INFINITY
80        };
81        entries[idx] = (idx as u8, scaled);
82    }
83    entries.sort_by(|a, b| b.1.total_cmp(&a.1));
84
85    let keep_k = if config.top_k == 0 {
86        entries.len()
87    } else {
88        config.top_k.min(entries.len())
89    };
90
91    let top_p = if config.top_p.is_finite() {
92        config.top_p.clamp(0.0, 1.0)
93    } else {
94        1.0
95    };
96
97    let mut max_logp = f64::NEG_INFINITY;
98    for &(_, logp) in entries.iter().take(keep_k) {
99        if logp.is_finite() {
100            max_logp = max_logp.max(logp);
101        }
102    }
103    if !max_logp.is_finite() {
104        return argmax_log_prob_byte(logps);
105    }
106
107    let mut weights = [(0u8, 0.0f64); 256];
108    let mut total = 0.0;
109    for (idx, &(byte, logp)) in entries.iter().take(keep_k).enumerate() {
110        let w = if logp.is_finite() {
111            (logp - max_logp).exp()
112        } else {
113            0.0
114        };
115        weights[idx] = (byte, w);
116        total += w;
117    }
118    if !(total.is_finite()) || total <= 0.0 {
119        return argmax_log_prob_byte(logps);
120    }
121
122    let cutoff_count = if top_p >= 1.0 {
123        keep_k
124    } else {
125        let mut cumulative = 0.0;
126        let mut keep = 0usize;
127        for &(_, w) in weights.iter().take(keep_k) {
128            cumulative += w / total;
129            keep += 1;
130            if cumulative >= top_p {
131                break;
132            }
133        }
134        keep.max(1)
135    };
136
137    let mut truncated_total = 0.0;
138    for &(_, w) in weights.iter().take(cutoff_count) {
139        truncated_total += w;
140    }
141    if !(truncated_total.is_finite()) || truncated_total <= 0.0 {
142        return argmax_log_prob_byte(logps);
143    }
144
145    let target = rng.next_f64() * truncated_total;
146    let mut cumulative = 0.0;
147    let mut picked = weights[0].0;
148    for &(byte, weight) in weights.iter().take(cutoff_count) {
149        cumulative += weight;
150        if cumulative >= target {
151            picked = byte;
152            break;
153        }
154    }
155    picked
156}
157
158pub(crate) fn try_generate_rate_backend_chain(
159    prefix_parts: &[&[u8]],
160    bytes: usize,
161    backend: &CompiledRateBackend,
162    config: GenerationConfig,
163) -> InfotheoryResult<Vec<u8>> {
164    if bytes == 0 {
165        return Ok(Vec::new());
166    }
167
168    let total = prefix_parts
169        .iter()
170        .map(|p| p.len() as u64)
171        .sum::<u64>()
172        .saturating_add(bytes as u64);
173    let mut session =
174        RateBackendSession::from_backend(backend.clone(), Some(total)).map_err(|e| {
175            InfotheoryError::runtime(format!("rate backend generation init failed: {e}"))
176        })?;
177    for &part in prefix_parts {
178        session.observe(part);
179    }
180    let out = session.generate_bytes(bytes, config);
181    session.finish().map_err(|e| {
182        InfotheoryError::runtime(format!("rate backend generation finalize failed: {e}"))
183    })?;
184    Ok(out)
185}
186
187#[cfg(test)]
188#[allow(dead_code)]
189pub(crate) fn generate_rate_backend_chain(
190    prefix_parts: &[&[u8]],
191    bytes: usize,
192    backend: &CompiledRateBackend,
193    config: GenerationConfig,
194) -> Vec<u8> {
195    try_generate_rate_backend_chain(prefix_parts, bytes, backend, config)
196        .expect("generate_rate_backend_chain")
197}
198
199/// Generate a continuation from `prompt`
200/// using the current default context and [`GenerationConfig::default()`].
201///
202/// The default is deterministic frozen sampling with seed `42`.
203#[inline(always)]
204pub fn try_generate_bytes(prompt: &[u8], bytes: usize) -> InfotheoryResult<Vec<u8>> {
205    with_default_ctx(|ctx| ctx.try_generate_bytes(prompt, bytes))
206}
207
208/// Generate a continuation from `prompt` using the current default context.
209#[inline(always)]
210pub fn try_generate_bytes_with_config(
211    prompt: &[u8],
212    bytes: usize,
213    config: GenerationConfig,
214) -> InfotheoryResult<Vec<u8>> {
215    with_default_ctx(|ctx| ctx.try_generate_bytes_with_config(prompt, bytes, config))
216}
217
218/// Generate a continuation after conditioning on an explicit chain of prefix parts
219/// using the current default context and [`GenerationConfig::default()`].
220#[inline(always)]
221pub fn try_generate_bytes_conditional_chain(
222    prefix_parts: &[&[u8]],
223    bytes: usize,
224) -> InfotheoryResult<Vec<u8>> {
225    with_default_ctx(|ctx| ctx.try_generate_bytes_conditional_chain(prefix_parts, bytes))
226}
227
228/// Generate a continuation after conditioning on an explicit chain of prefix parts
229/// using the current default context.
230#[inline(always)]
231pub fn try_generate_bytes_conditional_chain_with_config(
232    prefix_parts: &[&[u8]],
233    bytes: usize,
234    config: GenerationConfig,
235) -> InfotheoryResult<Vec<u8>> {
236    with_default_ctx(|ctx| {
237        ctx.try_generate_bytes_conditional_chain_with_config(prefix_parts, bytes, config)
238    })
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    fn sparse_logps(entries: &[(u8, f64)]) -> [f64; 256] {
246        let mut logps = [f64::NEG_INFINITY; 256];
247        for &(byte, logp) in entries {
248            logps[byte as usize] = logp;
249        }
250        logps
251    }
252
253    #[test]
254    fn greedy_generation_picks_argmax() {
255        let logps = sparse_logps(&[(7, -0.2), (42, -1.0)]);
256        let mut rng = GenerationRng::new(123);
257        let picked = pick_generated_byte(&logps, GenerationConfig::greedy_frozen(), &mut rng);
258        assert_eq!(picked, 7);
259    }
260
261    #[test]
262    fn nonpositive_temperature_falls_back_to_argmax() {
263        let logps = sparse_logps(&[(3, -0.1), (11, -0.3)]);
264        let mut rng = GenerationRng::new(7);
265        let mut config = GenerationConfig::sampled_frozen(7);
266        config.temperature = 0.0;
267        let picked = pick_generated_byte(&logps, config, &mut rng);
268        assert_eq!(picked, 3);
269    }
270
271    #[test]
272    fn top_k_sampling_respects_truncation() {
273        let logps = sparse_logps(&[(9, -0.01), (10, -0.02), (11, -0.03)]);
274        let mut rng = GenerationRng::new(99);
275        let mut config = GenerationConfig::sampled_frozen(99);
276        config.top_k = 1;
277        let picked = pick_generated_byte(&logps, config, &mut rng);
278        assert_eq!(picked, 9);
279    }
280
281    #[test]
282    fn top_p_sampling_keeps_only_minimal_prefix_mass() {
283        let logps = sparse_logps(&[(5, 0.0), (6, -1.5), (7, -3.0)]);
284        let mut rng = GenerationRng::new(5);
285        let mut config = GenerationConfig::sampled_frozen(5);
286        config.top_p = 0.5;
287        let picked = pick_generated_byte(&logps, config, &mut rng);
288        assert_eq!(picked, 5);
289    }
290
291    #[test]
292    fn nonfinite_log_probs_do_not_panic() {
293        let mut logps = [f64::NEG_INFINITY; 256];
294        logps[0] = f64::NAN;
295        let mut rng = GenerationRng::new(17);
296        let picked = pick_generated_byte(&logps, GenerationConfig::sampled_frozen(17), &mut rng);
297        assert_eq!(picked, 0);
298    }
299
300    #[test]
301    fn unit_interval_mapping_excludes_one() {
302        assert_eq!(unit_interval_from_u64(0), 0.0);
303        assert!(unit_interval_from_u64(u64::MAX) < 1.0);
304    }
305}
infotheory/api/generation.rs

infotheory/api/
generation.rs
