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

use std::alloc::System;

use stats_alloc::{Region, StatsAlloc, INSTRUMENTED_SYSTEM};
use token_sampling::{
    ChainConfig, Dist, Dry, MirostatV1, MirostatV2, SamplerChain, SplitMix64, Temperature,
};

#[global_allocator]
static GLOBAL: &StatsAlloc<System> = &INSTRUMENTED_SYSTEM;

#[test]
fn defect_d5_steady_state_sampling_allocates_zero_bytes() {
    for &size in &[4, 32, 1024, 131_072] {
        let logits = make_logits(size);
        let history = (0..size.min(96))
            .map(|value| value as u32)
            .collect::<Vec<_>>();
        check_chain(
            SamplerChain::llama_default(ChainConfig::default()),
            &logits,
            &history,
        );
        check_chain(
            SamplerChain::new()
                .push(Temperature { t: 0.8 })
                .push(MirostatV1::new(size as u32)),
            &logits,
            &history,
        );
        check_chain(
            SamplerChain::new()
                .push(Temperature { t: 0.8 })
                .push(MirostatV2::default()),
            &logits,
            &history,
        );
        check_chain(
            SamplerChain::new()
                .push(Dry {
                    multiplier: 0.8,
                    base: 1.75,
                    allowed_length: 2,
                    last_n: 64,
                    breakers: vec![vec![13], vec![27, 28]],
                })
                .push(Dist),
            &logits,
            &history,
        );
    }
}

fn check_chain(mut chain: SamplerChain, logits: &[f32], history: &[u32]) {
    let mut rng = SplitMix64::new(77);
    let mut warmup = logits.to_vec();
    chain.sample(&mut warmup, history, &mut rng).unwrap();
    let mut next = logits.to_vec();
    let region = Region::new(GLOBAL);
    chain.sample(&mut next, history, &mut rng).unwrap();
    let change = region.change();
    assert_eq!(
        change.allocations,
        0,
        "allocation count at vocab {}",
        logits.len()
    );
    assert_eq!(
        change.reallocations,
        0,
        "reallocation count at vocab {}",
        logits.len()
    );
    assert_eq!(
        change.bytes_allocated,
        0,
        "allocated bytes at vocab {}",
        logits.len()
    );
}

fn make_logits(size: usize) -> Vec<f32> {
    (0..size)
        .map(|index| {
            let value = (index as u64)
                .wrapping_mul(6_364_136_223_846_793_005_u64)
                .wrapping_add(1_442_695_040_888_963_407);
            ((value >> 40) as f32 / (1_u32 << 24) as f32) * 12.0 - 6.0
        })
        .collect()
}