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
41
42
43
44
45
46
47
48
49
50
51

use std::collections::{HashMap, HashSet};

use serde::Deserialize;

// ── Raw deserialization helper ───────────────────────────────────────────

#[derive(Deserialize)]
struct RawWordLevel {
    #[serde(default)]
    vocab: HashMap<String, u32>,
    #[serde(default)]
    unk_token: String,
}

// ── WordLevel model ──────────────────────────────────────────────────────

/// Whole-word lookup model: each pre-tokenized word maps to exactly one
/// token — itself when in the vocabulary, `unk_token` otherwise.
#[derive(Clone, Debug, Deserialize)]
#[serde(from = "RawWordLevel")]
pub struct WordLevel {
    /// Vocabulary keys; token ids are irrelevant for counting.
    vocab: HashSet<String>,
    /// Whether `unk_token` itself resolves to a vocabulary entry.
    unk_in_vocab: bool,
}

impl From<RawWordLevel> for WordLevel {
    fn from(raw: RawWordLevel) -> Self {
        let unk_in_vocab = raw.vocab.contains_key(&raw.unk_token);
        Self {
            vocab: raw.vocab.into_keys().collect(),
            unk_in_vocab,
        }
    }
}

impl WordLevel {
    /// Count tokens for one pre-tokenized word: always 1, unless the word
    /// is out of vocabulary and `unk_token` is missing from the vocabulary.
    #[inline]
    pub fn count_tokens(&self, input: &str) -> Result<usize, String> {
        if self.vocab.contains(input) || self.unk_in_vocab {
            Ok(1)
        } else {
            Err(format!(
                "WordLevel: `{input}` not in vocab and unk token is missing"
            ))
        }
    }
}