wordchipper 0.9.2

HPC Rust LLM Tokenizer Library
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

//! # Vocab Testing Tools

use crate::{
    TokenType,
    alloc::vec::Vec,
    spanners::TextSpanningConfig,
    vocab::{
        ByteMapVocab,
        SpanMapVocab,
        SpanTokenMap,
        UnifiedTokenVocab,
    },
};

/// Create a test [`UnifiedTokenVocab`].
pub fn build_test_vocab<T: TokenType, C>(
    byte_vocab: ByteMapVocab<T>,
    segmentation: C,
) -> UnifiedTokenVocab<T>
where
    C: Into<TextSpanningConfig<T>>,
{
    let mut span_map: SpanTokenMap<T> = Default::default();
    span_map.extend(
        [
            ("at", 300),
            ("ate", 301),
            ("th", 302),
            ("the", 303),
            (", ", 304),
            ("on", 305),
            ("he", 306),
            ("ll", 307),
            ("hell", 308),
            ("hello", 309),
            ("wo", 310),
            ("ld", 311),
            ("rld", 312),
            ("world", 313),
            ("fo", 314),
            ("for", 315),
            ("all", 316),
            (". ", 317),
        ]
        .into_iter()
        .map(|(k, v)| (k.as_bytes().to_vec(), T::from_usize(v).unwrap())),
    );

    let span_vocab = SpanMapVocab::new(byte_vocab, span_map).unwrap();

    UnifiedTokenVocab::from_span_vocab(segmentation.into(), span_vocab).unwrap()
}

/// Build a [`ByteMapVocab`] with all tokens shifted by `shift`.
///
/// This is a purposely stupid byte map; useful for testing.
pub fn build_test_shift_byte_vocab<T: TokenType>(shift: usize) -> ByteMapVocab<T> {
    // This is a purposely stupid byte map.
    ByteMapVocab::<T>::from_byte_to_token(
        &ByteMapVocab::<T>::default()
            .byte_tokens()
            .iter()
            .map(|&t| t + T::from_usize(shift).unwrap())
            .collect::<Vec<T>>(),
    )
}