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
use serde::Deserialize;
use crate::pre_tokenized::PreTokenizedString;
use crate::pre_tokenized::PtSplit;
use super::Error;
/// A compiled Digits pre-tokenizer.
///
/// Isolates digit sequences from surrounding text. When `individual_digits` is
/// `true`, each digit becomes a separate token.
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct Digits {
#[serde(default)]
individual_digits: bool,
}
impl Digits {
/// Refine the splits of a [`PreTokenizedString`] in place.
pub fn pre_tokenize(&self, pts: &mut PreTokenizedString) -> Result<(), Error> {
let buffer = pts.buffer.as_str();
let old_splits = &pts.splits;
let hint: usize = old_splits.len() << 1; // Vec::len <= isize::MAX
let mut new_splits = Vec::with_capacity(hint);
for split in old_splits {
if split.token_id.is_some() {
new_splits.push(split.clone());
continue;
}
let text = &buffer[split.range.clone()];
if text.is_empty() {
continue;
}
let base = split.range.start;
if self.individual_digits {
// Each digit gets its own segment; non-digits are grouped.
let mut current_start = 0;
let mut current_is_digit = text.chars().next().is_some_and(|c| c.is_ascii_digit());
for (i, ch) in text.char_indices() {
let is_digit = ch.is_ascii_digit();
if is_digit != current_is_digit {
if i > current_start {
new_splits.push(PtSplit {
range: (base + current_start)..(base + i),
token_id: None,
});
}
current_start = i;
current_is_digit = is_digit;
}
}
if current_start < text.len() {
new_splits.push(PtSplit {
range: (base + current_start)..(base + text.len()),
token_id: None,
});
}
} else {
// Group consecutive digits together; non-digits are separate.
let mut current_start = 0;
let mut current_is_digit = text.chars().next().is_some_and(|c| c.is_ascii_digit());
let mut any_digits = current_is_digit;
for (i, ch) in text.char_indices() {
let is_digit = ch.is_ascii_digit();
any_digits = any_digits || is_digit;
if is_digit != current_is_digit {
if i > current_start {
new_splits.push(PtSplit {
range: (base + current_start)..(base + i),
token_id: None,
});
}
current_start = i;
current_is_digit = is_digit;
}
}
if current_start < text.len() {
new_splits.push(PtSplit {
range: (base + current_start)..(base + text.len()),
token_id: None,
});
}
// If there were no digits, restore the original split.
if !any_digits {
new_splits.truncate(
new_splits
.len()
.saturating_sub(text.chars().filter(|c| !c.is_ascii_digit()).count()),
);
new_splits.push(split.clone());
}
}
}
pts.splits = new_splits;
Ok(())
}
}