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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

use std::ops::Range;

/// A split within a [`PreTokenizedString`]'s buffer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PtSplit {
    /// Byte range into the parent buffer.
    pub range: Range<usize>,
    /// If `Some`, this split is an added token and should count as 1.
    pub token_id: Option<u32>,
}

impl PtSplit {
    #[inline]
    pub fn from_range(range: Range<usize>) -> Self {
        Self {
            range,
            token_id: None,
        }
    }

    #[inline]
    pub fn from_token(span: usize, token: u32) -> Self {
        Self {
            range: span..span,
            token_id: Some(token),
        }
    }
}

/// A single-buffer representation of pre-tokenized text.
#[derive(Debug, Clone)]
pub struct PreTokenizedString {
    pub buffer: String,
    pub splits: Vec<PtSplit>,
}

impl PreTokenizedString {
    /// An empty PreTokenizedString;
    pub const EMPTY: Self = Self {
        buffer: String::new(),
        splits: Vec::new(),
    };

    /// Create with a pre-built buffer and splits.
    pub fn new(buffer: String, splits: Vec<PtSplit>) -> Self {
        Self { buffer, splits }
    }

    /// Text content of a split.
    pub fn split_text(&self, split: &PtSplit) -> &str {
        &self.buffer[split.range.clone()]
    }

    /// Iterate over all non-added-token text spans.
    pub fn texts(&self) -> impl Iterator<Item = &str> + '_ {
        self.splits.iter().filter_map(|s| {
            if s.token_id.is_some() {
                return None;
            }
            let text = &self.buffer[s.range.clone()];
            if text.is_empty() { None } else { Some(text) }
        })
    }

    /// Number of splits that are pre-assigned added tokens (count as 1 each).
    pub fn added_token_count(&self) -> usize {
        self.splits.iter().filter(|s| s.token_id.is_some()).count()
    }
}