splintr 0.10.0

Fast Rust tokenizer (BPE + SentencePiece + WordPiece) with Python bindings
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

//! SentencePiece `Precompiled` charsmap normalizer.
//!
//! SentencePiece ships its normalization rules as a precompiled blob: a
//! darts-clone double-array trie mapping input byte sequences to offsets into a
//! concatenated, null-terminated table of normalized replacement strings. This
//! is what HuggingFace stores as `normalizer.precompiled_charsmap`.
//!
//! Matching it exactly (rather than approximating with NFKC) is required for
//! byte-identical parity on the long tail of Unicode (CJK-compatibility,
//! astral-plane, and other characters where the charsmap and NFKC disagree).

/// A parsed SentencePiece precompiled charsmap.
#[derive(Debug, Clone)]
pub struct Precompiled {
    /// darts-clone double-array units.
    trie: Vec<u32>,
    /// Null-terminated normalized replacement strings, concatenated.
    normalized: Vec<u8>,
}

// darts-clone unit accessors.
#[inline]
fn has_leaf(unit: u32) -> bool {
    ((unit >> 8) & 1) == 1
}
#[inline]
fn value(unit: u32) -> u32 {
    unit & 0x7fff_ffff
}
#[inline]
fn label(unit: u32) -> u32 {
    unit & 0x8000_00ff
}
#[inline]
fn offset(unit: u32) -> u32 {
    (unit >> 10) << ((unit & (1 << 9)) >> 6)
}

impl Precompiled {
    /// Parse a charsmap blob (the raw bytes of `precompiled_charsmap`).
    ///
    /// Layout: `u32 LE trie_size` · `trie_size` bytes of `u32 LE` trie units ·
    /// remaining bytes = the normalized-strings table.
    pub fn from_bytes(blob: &[u8]) -> Option<Self> {
        if blob.len() < 4 {
            return None;
        }
        let trie_size = u32::from_le_bytes([blob[0], blob[1], blob[2], blob[3]]) as usize;
        let trie_end = 4 + trie_size;
        if trie_end > blob.len() || !trie_size.is_multiple_of(4) {
            return None;
        }
        let trie = blob[4..trie_end]
            .chunks_exact(4)
            .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();
        let normalized = blob[trie_end..].to_vec();
        Some(Self { trie, normalized })
    }

    /// Longest common-prefix match of `key` in the trie, returning the matched
    /// byte length and the value (offset into `normalized`).
    fn longest_prefix(&self, key: &[u8]) -> Option<(usize, u32)> {
        let mut node_pos = 0usize;
        let mut unit = self.trie[0];
        node_pos ^= offset(unit) as usize;
        let mut best: Option<(usize, u32)> = None;
        for (i, &b) in key.iter().enumerate() {
            node_pos ^= b as usize;
            if node_pos >= self.trie.len() {
                break;
            }
            unit = self.trie[node_pos];
            if label(unit) != b as u32 {
                break;
            }
            node_pos ^= offset(unit) as usize;
            if has_leaf(unit) {
                if let Some(&v) = self.trie.get(node_pos) {
                    best = Some((i + 1, value(v)));
                }
            }
        }
        best
    }

    /// The null-terminated normalized string at `offset`.
    fn normalized_at(&self, offset: u32) -> &[u8] {
        let start = offset as usize;
        let mut end = start;
        while end < self.normalized.len() && self.normalized[end] != 0 {
            end += 1;
        }
        &self.normalized[start..end]
    }

    /// Normalize `text` by repeatedly replacing the longest matching prefix with
    /// its normalized form (SentencePiece's algorithm).
    pub fn normalize(&self, text: &str) -> String {
        let bytes = text.as_bytes();
        let mut out: Vec<u8> = Vec::with_capacity(bytes.len());
        let mut i = 0;
        while i < bytes.len() {
            match self.longest_prefix(&bytes[i..]) {
                Some((len, off)) if len > 0 => {
                    out.extend_from_slice(self.normalized_at(off));
                    i += len;
                }
                _ => {
                    // No rule: copy one whole UTF-8 character verbatim.
                    let ch_len = utf8_len(bytes[i]);
                    let end = (i + ch_len).min(bytes.len());
                    out.extend_from_slice(&bytes[i..end]);
                    i = end;
                }
            }
        }
        // The table is valid UTF-8 by construction; fall back lossily if not.
        String::from_utf8(out)
            .unwrap_or_else(|e| String::from_utf8_lossy(e.as_bytes()).into_owned())
    }
}

/// Byte length of the UTF-8 character starting with `first`.
#[inline]
fn utf8_len(first: u8) -> usize {
    if first < 0x80 {
        1
    } else if first >> 5 == 0b110 {
        2
    } else if first >> 4 == 0b1110 {
        3
    } else if first >> 3 == 0b11110 {
        4
    } else {
        1
    }
}