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
use std::ops::Range;
/// A split within a [`PreTokenizedString`]'s buffer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Split {
/// 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>,
}
/// A single-buffer representation of pre-tokenized text.
#[derive(Debug, Clone)]
pub struct PreTokenizedString {
buffer: String,
splits: Vec<Split>,
}
impl PreTokenizedString {
/// Create from a single text span (no pre-assigned tokens).
pub fn from_text(text: &str) -> Self {
let splits = if text.is_empty() {
Vec::new()
} else {
vec![Split {
range: 0..text.len(),
token_id: None,
}]
};
Self {
buffer: text.to_string(),
splits,
}
}
/// Create with a pre-built buffer and splits.
pub fn new(buffer: String, splits: Vec<Split>) -> Self {
Self { buffer, splits }
}
/// The underlying buffer.
pub fn buffer(&self) -> &str {
&self.buffer
}
/// The current splits.
pub fn splits(&self) -> &[Split] {
&self.splits
}
/// Text content of a split.
pub fn split_text(&self, split: &Split) -> &str {
&self.buffer[split.range.clone()]
}
/// Replace the buffer and splits entirely.
pub fn set_buffer(&mut self, buffer: String, splits: Vec<Split>) {
self.buffer = buffer;
self.splits = splits;
}
/// Replace only the splits, keeping the buffer unchanged.
pub fn refine_splits(&mut self, splits: Vec<Split>) {
self.splits = splits;
}
/// Count tokens for all splits sequentially.
///
/// For each text split, calls `count_fn` which increments the counter.
/// Added-token splits count as 1 each.
pub fn count_tokens<F>(&self, count_fn: F) -> Result<usize, String>
where
F: Fn(&str, &mut usize) -> Result<(), String>,
{
let mut count = 0usize;
for split in &self.splits {
if split.token_id.is_some() {
count += 1;
} else {
let text = self.split_text(split);
if !text.is_empty() {
count_fn(text, &mut count)?;
}
}
}
Ok(count)
}
/// Batched count: the callback receives the full buffer and all splits.
pub fn count_tokens_batched<F>(&self, count_fn: F) -> Result<usize, String>
where
F: Fn(&str, &[Split], &mut usize) -> Result<(), String>,
{
let mut count = 0usize;
count_fn(&self.buffer, &self.splits, &mut count)?;
Ok(count)
}
}