toklen 0.2.0

A single-threaded, lightweight, and fast token counter.
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

use std::borrow::Cow;

use crate::byte_level::BYTE_TO_UTF8;
use crate::byte_level::BYTE_TO_UTF8_LEN;

/// ByteLevel normalizer: replaces each byte of the UTF-8 input with its
/// GPT-2 byte-to-unicode mapping character.
///
/// This is the normalizer variant (no regex splitting or prefix-space logic —
/// those belong to the ByteLevel pre-tokenizer).
///
/// Returns `Cow::Borrowed` when the input is ASCII-only printable text
/// (bytes 0x21–0x7E, 0xA1–0xAC, 0xAE–0xFF, which map to themselves).
#[inline]
pub fn normalize(input: &str) -> Cow<'_, str> {
    let bytes = input.as_bytes();

    #[inline]
    fn check_byte(b: &u8) -> bool {
        !((b'!'..=b'~').contains(b) || (0xA1..=0xAC).contains(b) || *b >= 0xAE)
    }

    // Fast path: check if any byte needs remapping.
    let needs_encode = bytes.iter().any(check_byte);

    if !needs_encode {
        return Cow::Borrowed(input);
    }

    // Re-encode: each input byte → 1 or 2 UTF-8 bytes.
    let mut out = Vec::with_capacity(bytes.len() << 1);
    for &b in bytes {
        let utf8 = BYTE_TO_UTF8[b as usize];
        let len = BYTE_TO_UTF8_LEN[b as usize] as usize;
        out.extend_from_slice(&utf8[..len]);
    }

    // SAFETY: BYTE_TO_UTF8 contains only valid UTF-8 sequences.
    Cow::Owned(unsafe { String::from_utf8_unchecked(out) })
}