token-sampling 0.1.0

Composable token sampling over f32 logits with caller-supplied randomness
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

#[cfg(not(debug_assertions))]
use std::time::Instant;
#[cfg(not(debug_assertions))]
use std::{io, io::Write};

#[cfg(not(debug_assertions))]
use token_sampling::{ChainConfig, SamplerChain, SplitMix64};

#[cfg(not(debug_assertions))]
#[test]
fn default_chain_mean_is_below_two_milliseconds() {
    const SIZE: usize = 131_072;
    const ITERATIONS: u32 = 32;
    let mut logits = (0..SIZE)
        .map(|index| {
            let value = (index as u64)
                .wrapping_mul(2_862_933_555_777_941_757_u64)
                .wrapping_add(3_037_000_493);
            ((value >> 40) as f32 / (1_u32 << 24) as f32) * 12.0 - 6.0
        })
        .collect::<Vec<_>>();
    let mut chain = SamplerChain::llama_default(ChainConfig::default());
    let mut rng = SplitMix64::new(91);
    chain.sample(&mut logits, &[], &mut rng).unwrap();

    let start = Instant::now();
    for _ in 0..ITERATIONS {
        chain.sample(&mut logits, &[], &mut rng).unwrap();
    }
    let mean = start.elapsed().as_secs_f64() / f64::from(ITERATIONS);
    writeln!(
        io::stdout().lock(),
        "default chain mean: {:.3} ms",
        mean * 1000.0
    )
    .unwrap();
    assert!(mean < 0.002, "mean step took {:.3} ms", mean * 1000.0);
}