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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

use token_sampling::{
    ChainConfig, Dry, MinP, Penalties, SamplerChain, SplitMix64, TailFree, Temperature, TopK, TopP,
    Typical, Xtc,
};

#[test]
fn arbitrary_inputs_finish_with_a_token_or_typed_error() {
    let mut state = 0x243f_6a88_85a3_08d3_u64;
    for case in 0..=512 {
        let size = if case == 512 { 4096 } else { case % 257 };
        let mut logits = (0..size)
            .map(|index| match (case + index) % 41 {
                0 => f32::NAN,
                1 => f32::INFINITY,
                2 => f32::NEG_INFINITY,
                _ => {
                    state = state
                        .wrapping_mul(6_364_136_223_846_793_005)
                        .wrapping_add(1);
                    ((state >> 40) as f32 / (1_u32 << 24) as f32) * 40.0 - 20.0
                }
            })
            .collect::<Vec<_>>();
        let history = (0..case % 384)
            .map(|value| (value % 311) as u32)
            .collect::<Vec<_>>();
        let config = ChainConfig {
            penalties: Penalties {
                last_n: case as i32 - 100,
                repeat: (case as f32 - 200.0) / 50.0,
                freq: (case as f32 - 256.0) / 30.0,
                presence: (case as f32 - 128.0) / 20.0,
            },
            dry: Dry {
                multiplier: 0.5 + (case % 13) as f32 / 10.0,
                base: 1.0 + (case % 9) as f32 / 20.0,
                allowed_length: (case % 7) as i32,
                last_n: 8 + (case % 96) as i32,
                breakers: vec![vec![], vec![1], vec![2, 3]],
            },
            top_k: TopK {
                k: case as i32 - 64,
            },
            tail_free: TailFree {
                z: (case as f32 - 128.0) / 128.0,
                min_keep: case % 8,
            },
            typical: Typical {
                p: (case as f32 - 128.0) / 128.0,
                min_keep: case % 8,
            },
            top_p: TopP {
                p: (case as f32 - 128.0) / 128.0,
                min_keep: case % 8,
            },
            min_p: MinP {
                p: (case as f32 - 128.0) / 128.0,
                min_keep: case % 8,
            },
            xtc: Xtc {
                probability: (case as f32 - 128.0) / 128.0,
                threshold: (case as f32 - 64.0) / 128.0,
                min_keep: case % 8,
            },
            temperature: token_sampling::DynTemperature {
                t: (case as f32 - 100.0) / 50.0,
                delta: (case as f32 - 200.0) / 100.0,
                exponent: (case as f32 - 128.0) / 64.0,
            },
        };
        let mut chain = SamplerChain::llama_default(config);
        let mut rng = SplitMix64::new(case as u64);
        let result = chain.sample(&mut logits, &history, &mut rng);
        if let Ok(token) = result {
            assert!((token as usize) < size);
        }
    }
}

#[test]
fn non_finite_stage_parameters_do_not_panic() {
    let mut chain = SamplerChain::new()
        .push(Temperature { t: f32::NAN })
        .push(TopP {
            p: f32::NAN,
            min_keep: usize::MAX,
        })
        .push(Xtc {
            probability: f32::NAN,
            threshold: f32::NAN,
            min_keep: usize::MAX,
        })
        .push(token_sampling::Dist);
    let mut logits = [0.0, 1.0];
    let mut rng = SplitMix64::new(1);
    assert!(chain.sample(&mut logits, &[], &mut rng).is_ok());
}