logicaffeine_compile/model/tokenizer/
split.rs

1//! Pre-tokenization: a fragment of text split into words before merging, as llama.cpp's `unicode_regex_split` does.
2//!
3//! The patterns apply in order, each refining the spans the last one left (spans count codepoints). A pattern
4//! llama.cpp has a hand-written splitter for runs that splitter ([`Custom`]); every other runs through
5//! [`super::regex`], over the *collapsed* text when it names a Unicode category (`\p{L}` …, each non-ASCII codepoint
6//! folded to one byte standing for its category) and over the codepoints otherwise. A word of a byte-encoded
7//! vocabulary is then written in GPT-2's byte alphabet.
8
9use super::regex::{Regex, RegexError, Width};
10use super::unicode::{self, cpts_from_utf8, flags, flags_in, is_han, push_utf8, to_lower, Flags, InvalidCodepoint};
11
12const OUT_OF_RANGE: u32 = 0xFFFF_FFFF;
13
14/// The hand-written splitters, each standing for the pattern llama.cpp dispatches to it.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub enum Custom {
17    Gpt2,
18    Llama3,
19    Qwen2,
20    Qwen35,
21    K2Horizon,
22    KimiK2,
23    Afmoe,
24    Newlines,
25}
26
27impl Custom {
28    /// The splitter llama.cpp's `unicode_regex_split_custom` runs for `pattern`, if any.
29    pub fn for_pattern(pattern: &str) -> Option<Custom> {
30        Some(match pattern {
31            "'s|'t|'re|'ve|'m|'ll|'d| ?\\p{L}+| ?\\p{N}+| ?[^\\s\\p{L}\\p{N}]+|\\s+(?!\\S)" => Custom::Gpt2,
32            "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
33            | "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" => Custom::Llama3,
34            "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" => Custom::Qwen2,
35            "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" => Custom::Qwen35,
36            "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}|\\p{M}|\\u200C|\\u200D)+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" => Custom::K2Horizon,
37            "\\p{Han}+" => Custom::KimiK2,
38            "\\p{AFMoE_digits}" | "\\d{1,3}(?=(?:\\d{3})*\\b)" => Custom::Afmoe,
39            "[^\\n]+|[\\n]+" => Custom::Newlines,
40            _ => return None,
41        })
42    }
43
44    fn split(self, cpts: &[u32], offsets: &[usize]) -> Vec<usize> {
45        match self {
46            Custom::Gpt2 => gpt2(cpts, offsets),
47            Custom::Llama3 => llama3_family(cpts, offsets, Family::LLAMA3),
48            Custom::Qwen2 => llama3_family(cpts, offsets, Family::QWEN2),
49            Custom::Qwen35 => llama3_family(cpts, offsets, Family::QWEN35),
50            Custom::K2Horizon => llama3_family(cpts, offsets, Family::K2_HORIZON),
51            Custom::KimiK2 => kimi_k2(cpts, offsets),
52            Custom::Afmoe => afmoe(cpts, offsets),
53            Custom::Newlines => newlines(cpts, offsets),
54        }
55    }
56}
57
58/// The Unicode category escapes the collapsed text stands in for.
59const CATEGORIES: [(&str, u16); 10] = [
60    ("\\p{N}", Flags::NUMBER),
61    ("\\p{L}", Flags::LETTER),
62    ("\\p{P}", Flags::PUNCTUATION),
63    ("\\p{M}", Flags::ACCENT_MARK),
64    ("\\p{S}", Flags::SYMBOL),
65    ("\\p{Lu}", Flags::LETTER),
66    ("\\p{Ll}", Flags::LETTER),
67    ("\\p{Lt}", Flags::LETTER),
68    ("\\p{Lm}", Flags::LETTER),
69    ("\\p{Lo}", Flags::LETTER),
70];
71
72/// The byte a category collapses to, and the ASCII members its class keeps (llama.cpp's `k_ucat_cpt`, `k_ucat_map`).
73fn collapsed_category(cat: u16) -> Option<(u8, &'static [u8])> {
74    Some(match cat {
75        Flags::NUMBER => (0xD1, b"0-9"),
76        Flags::LETTER => (0xD2, b"A-Za-z"),
77        Flags::PUNCTUATION => (0xD3, b"!-#%-*,-/:-;?-@\\[-\\]_\\{\\}"),
78        Flags::ACCENT_MARK => (0xD4, b""),
79        Flags::SYMBOL => (0xD5, b"\\$\\+<->^`\\|\\~"),
80        _ => return None,
81    })
82}
83
84fn names_category(pattern: &str) -> bool {
85    CATEGORIES.iter().any(|(k, _)| pattern.contains(k))
86}
87
88/// A pattern naming categories, rewritten over the collapsed alphabet.
89fn collapse_pattern(pattern: &str) -> Result<Vec<u32>, RegexError> {
90    let re = pattern.as_bytes();
91    if !pattern.is_ascii() {
92        return Err(RegexError {
93            pattern: pattern.to_string(),
94            at: 0,
95            why: "a pattern names a Unicode category and holds a non-ASCII character",
96        });
97    }
98    let mut out: Vec<u8> = Vec::with_capacity(re.len() * 2);
99    let mut inside = false;
100    let mut i = 0;
101    while i < re.len() {
102        if re[i] == b'[' && (i == 0 || re[i - 1] != b'\\') {
103            out.push(b'[');
104            inside = true;
105            i += 1;
106            continue;
107        }
108        if inside && re[i] == b']' && re[i - 1] != b'\\' {
109            out.push(b']');
110            inside = false;
111            i += 1;
112            continue;
113        }
114        if re[i] == b'\\' && i + 3 < re.len() && re[i + 1] == b'p' && re[i + 2] == b'{' {
115            if let Some(close) = re[i + 3..].iter().position(|&c| c == b'}').map(|k| k + i + 3) {
116                if close <= i + 10 {
117                    let name = &pattern[i..=close];
118                    if let Some(&(_, cat)) = CATEGORIES.iter().find(|(k, _)| *k == name) {
119                        let (byte, ascii) = collapsed_category(cat).expect("every named category collapses");
120                        if !inside {
121                            out.push(b'[');
122                        }
123                        out.push(byte);
124                        out.extend_from_slice(ascii);
125                        if !inside {
126                            out.push(b']');
127                        }
128                        i = close + 1;
129                        continue;
130                    }
131                }
132            }
133        }
134        out.push(re[i]);
135        i += 1;
136    }
137    Ok(out.into_iter().map(u32::from).collect())
138}
139
140#[derive(Clone, Debug)]
141enum Step {
142    Custom(Custom),
143    Collapsed(Regex),
144    Wide(Regex),
145}
146
147/// A compiled pre-tokenizer: its patterns, in order, and whether words are byte-encoded.
148#[derive(Clone, Debug)]
149pub struct Splitter {
150    steps: Vec<Step>,
151    collapse: bool,
152    byte_encode: bool,
153}
154
155impl Splitter {
156    /// Compile the patterns of a pre-tokenizer; a pattern the regular-expression engine refuses is named.
157    pub fn new(patterns: &[&str], byte_encode: bool) -> Result<Splitter, RegexError> {
158        let mut steps = Vec::with_capacity(patterns.len());
159        for p in patterns {
160            steps.push(match Custom::for_pattern(p) {
161                Some(c) => Step::Custom(c),
162                None if names_category(p) => {
163                    Step::Collapsed(Regex::from_units(&collapse_pattern(p)?, p, Width::Narrow)?)
164                }
165                None => Step::Wide(Regex::new(p, Width::Wide)?),
166            });
167        }
168        Ok(Splitter { steps, collapse: patterns.iter().any(|p| names_category(p)), byte_encode })
169    }
170
171    /// The words `text` splits into, each byte-encoded when the vocabulary is.
172    pub fn split(&self, text: &[u8]) -> Result<Vec<Vec<u8>>, InvalidCodepoint> {
173        if text.is_empty() {
174            return Ok(Vec::new());
175        }
176        let cpts = cpts_from_utf8(text);
177        let collapsed: Vec<u32> = if self.collapse {
178            cpts.iter()
179                .map(|&c| {
180                    if c < 128 {
181                        return c;
182                    }
183                    let f = flags(c);
184                    if f.is_whitespace() {
185                        0x0B
186                    } else {
187                        collapsed_category(f.category()).map_or(0xD0, |(b, _)| b as u32)
188                    }
189                })
190                .collect()
191        } else {
192            Vec::new()
193        };
194        let mut offsets = vec![cpts.len()];
195        for step in &self.steps {
196            offsets = match step {
197                Step::Custom(c) => c.split(&cpts, &offsets),
198                Step::Collapsed(re) => stl_split(re, &collapsed, &offsets),
199                Step::Wide(re) => {
200                    let wide: Vec<u32> =
201                        cpts.iter().map(|&c| if c > 0x7F && flags(c).is_whitespace() { 0x0B } else { c }).collect();
202                    stl_split(re, &wide, &offsets)
203                }
204            };
205        }
206        let mut words = Vec::with_capacity(offsets.len());
207        let mut start = 0;
208        for len in offsets {
209            let mut w = Vec::new();
210            for &c in &cpts[start..start + len] {
211                push_utf8(c, &mut w)?;
212            }
213            start += len;
214            if self.byte_encode {
215                let mut enc = Vec::with_capacity(w.len() * 2);
216                for b in w {
217                    unicode::byte_to_utf8(b, &mut enc);
218                }
219                w = enc;
220            }
221            words.push(w);
222        }
223        Ok(words)
224    }
225}
226
227/// `unicode_regex_split_stl`: each span split at the pattern's matches, the text between matches kept as words.
228fn stl_split(re: &Regex, units: &[u32], offsets: &[usize]) -> Vec<usize> {
229    let mut out = Vec::with_capacity(offsets.len());
230    let mut start = 0;
231    for &offset in offsets {
232        let span = &units[start..start + offset];
233        let mut at = 0;
234        for (pos, len) in re.matches(span) {
235            if pos > at {
236                out.push(pos - at);
237            }
238            out.push(len);
239            at = pos + len;
240        }
241        if at < offset {
242            out.push(offset - at);
243        }
244        start += offset;
245    }
246    out
247}
248
249/// One span of the hand-written splitters: its bounds over `cpts`, and the words emitted so far.
250struct Span<'a> {
251    cpts: &'a [u32],
252    lo: usize,
253    hi: usize,
254    prev_end: usize,
255}
256
257impl Span<'_> {
258    fn cpt(&self, pos: usize) -> u32 {
259        if self.lo <= pos && pos < self.hi {
260            self.cpts[pos]
261        } else {
262            OUT_OF_RANGE
263        }
264    }
265
266    fn flags(&self, pos: usize) -> Flags {
267        flags_in(self.cpts, self.lo, self.hi, pos)
268    }
269
270    /// Emit the word ending at `end` (nothing when it is empty); its length.
271    fn add(&mut self, end: usize, out: &mut Vec<usize>) -> usize {
272        let len = end - self.prev_end;
273        if len > 0 {
274            out.push(len);
275        }
276        self.prev_end = end;
277        len
278    }
279}
280
281fn spans<'a>(cpts: &'a [u32], offsets: &'a [usize]) -> impl Iterator<Item = Span<'a>> + 'a {
282    let mut start = 0;
283    offsets.iter().map(move |&offset| {
284        let s = Span { cpts, lo: start, hi: start + offset, prev_end: start };
285        start += offset;
286        s
287    })
288}
289
290/// GPT-2: `'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`.
291fn gpt2(cpts: &[u32], offsets: &[usize]) -> Vec<usize> {
292    let mut out = Vec::with_capacity(offsets.len());
293    for mut s in spans(cpts, offsets) {
294        let mut pos = s.lo;
295        while pos < s.hi {
296            let cpt = s.cpt(pos);
297            let flags = s.flags(pos);
298            if cpt == '\'' as u32 && pos + 1 < s.hi {
299                let next = s.cpt(pos + 1);
300                if [b's', b't', b'm', b'd'].iter().any(|&c| next == c as u32) {
301                    pos += s.add(pos + 2, &mut out);
302                    continue;
303                }
304                if pos + 2 < s.hi {
305                    let nn = s.cpt(pos + 2);
306                    if (next == 'r' as u32 && nn == 'e' as u32)
307                        || (next == 'v' as u32 && nn == 'e' as u32)
308                        || (next == 'l' as u32 && nn == 'l' as u32)
309                    {
310                        pos += s.add(pos + 3, &mut out);
311                        continue;
312                    }
313                }
314            }
315            let mut flags2 = if cpt == ' ' as u32 { s.flags(pos + 1) } else { flags };
316            if flags2.is_letter() {
317                pos += (cpt == ' ' as u32) as usize;
318                while flags2.is_letter() {
319                    pos += 1;
320                    flags2 = s.flags(pos);
321                }
322                s.add(pos, &mut out);
323                continue;
324            }
325            if flags2.is_number() {
326                pos += (cpt == ' ' as u32) as usize;
327                while flags2.is_number() {
328                    pos += 1;
329                    flags2 = s.flags(pos);
330                }
331                s.add(pos, &mut out);
332                continue;
333            }
334            let other = |f: Flags| !(f.is_whitespace() || f.is_letter() || f.is_number()) && f.0 != 0;
335            if other(flags2) {
336                pos += (cpt == ' ' as u32) as usize;
337                while other(flags2) {
338                    pos += 1;
339                    flags2 = s.flags(pos);
340                }
341                s.add(pos, &mut out);
342                continue;
343            }
344            let mut ws = 0;
345            while s.flags(pos + ws).is_whitespace() {
346                ws += 1;
347            }
348            if ws > 1 && s.cpt(pos + ws) != OUT_OF_RANGE {
349                pos += ws - 1;
350                s.add(pos, &mut out);
351                continue;
352            }
353            if ws > 0 {
354                pos += ws;
355                s.add(pos, &mut out);
356                continue;
357            }
358            pos += 1;
359            s.add(pos, &mut out);
360        }
361    }
362    out
363}
364
365#[derive(Clone, Copy, PartialEq, Eq)]
366enum WordRun {
367    /// `\p{L}+`
368    Letters,
369    /// `[\p{L}\p{M}]+`
370    LettersMarks,
371    /// `(?:\p{L}|\p{M}|‌|‍)+`
372    LettersMarksJoiners,
373}
374
375/// The splitters llama.cpp writes for the Llama 3 pattern and its relatives: they differ in what a word run holds,
376/// whether digits group by three, whether a run of punctuation stops at marks, and K2-Horizon's long-s fold.
377#[derive(Clone, Copy)]
378struct Family {
379    word: WordRun,
380    digits_by_three: bool,
381    punct_stops_at_marks: bool,
382    long_s: bool,
383}
384
385impl Family {
386    const LLAMA3: Family =
387        Family { word: WordRun::Letters, digits_by_three: true, punct_stops_at_marks: false, long_s: false };
388    const QWEN2: Family =
389        Family { word: WordRun::Letters, digits_by_three: false, punct_stops_at_marks: false, long_s: false };
390    const QWEN35: Family =
391        Family { word: WordRun::LettersMarks, digits_by_three: false, punct_stops_at_marks: true, long_s: false };
392    const K2_HORIZON: Family =
393        Family { word: WordRun::LettersMarksJoiners, digits_by_three: true, punct_stops_at_marks: false, long_s: true };
394}
395
396/// `(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?<word>+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`
397/// and its relatives.
398fn llama3_family(cpts: &[u32], offsets: &[usize], fam: Family) -> Vec<usize> {
399    let mut out = Vec::with_capacity(offsets.len());
400    for mut s in spans(cpts, offsets) {
401        let word = |s: &Span, p: usize| -> bool {
402            let f = s.flags(p);
403            match fam.word {
404                WordRun::Letters => f.is_letter(),
405                WordRun::LettersMarks => f.is_letter() || f.is_accent_mark(),
406                WordRun::LettersMarksJoiners => {
407                    let c = s.cpt(p);
408                    c == 0x200C || c == 0x200D || f.is_letter() || f.is_accent_mark()
409                }
410            }
411        };
412        let mut pos = s.lo;
413        while pos < s.hi {
414            let cpt = s.cpt(pos);
415            let flags = s.flags(pos);
416            if cpt == '\'' as u32 && pos + 1 < s.hi {
417                let mut next = to_lower(s.cpt(pos + 1));
418                if fam.long_s && next == 0x017F {
419                    next = 's' as u32;
420                }
421                if [b's', b't', b'm', b'd'].iter().any(|&c| next == c as u32) {
422                    pos += s.add(pos + 2, &mut out);
423                    continue;
424                }
425                if pos + 2 < s.hi {
426                    let nn = to_lower(s.cpt(pos + 2));
427                    if (next == 'r' as u32 && nn == 'e' as u32)
428                        || (next == 'v' as u32 && nn == 'e' as u32)
429                        || (next == 'l' as u32 && nn == 'l' as u32)
430                    {
431                        pos += s.add(pos + 3, &mut out);
432                        continue;
433                    }
434                }
435            }
436            if !(cpt == '\r' as u32 || cpt == '\n' as u32 || flags.is_number()) && (word(&s, pos) || word(&s, pos + 1)) {
437                pos += 1;
438                while word(&s, pos) {
439                    pos += 1;
440                }
441                s.add(pos, &mut out);
442                continue;
443            }
444            if flags.is_number() {
445                if fam.digits_by_three {
446                    let mut ini = pos;
447                    while s.flags(pos).is_number() {
448                        pos += 1;
449                        if pos - ini >= 3 {
450                            s.add(pos, &mut out);
451                            ini = pos;
452                        }
453                    }
454                } else {
455                    pos += 1;
456                }
457                s.add(pos, &mut out);
458                continue;
459            }
460            let excluded = |f: Flags| {
461                f.is_whitespace() || f.is_letter() || f.is_number() || (fam.punct_stops_at_marks && f.is_accent_mark())
462            };
463            let mut flags2 = if cpt == ' ' as u32 { s.flags(pos + 1) } else { flags };
464            if !excluded(flags2) && flags.0 != 0 {
465                pos += (cpt == ' ' as u32) as usize;
466                while !excluded(flags2) && flags2.0 != 0 {
467                    pos += 1;
468                    flags2 = s.flags(pos);
469                }
470                let mut c2 = s.cpt(pos);
471                while c2 == '\r' as u32 || c2 == '\n' as u32 {
472                    pos += 1;
473                    c2 = s.cpt(pos);
474                }
475                s.add(pos, &mut out);
476                continue;
477            }
478            let mut ws = 0;
479            let mut last_rn = 0;
480            while s.flags(pos + ws).is_whitespace() {
481                let c2 = s.cpt(pos + ws);
482                if c2 == '\r' as u32 || c2 == '\n' as u32 {
483                    last_rn = pos + ws + 1;
484                }
485                ws += 1;
486            }
487            if last_rn > 0 {
488                pos = last_rn;
489                s.add(pos, &mut out);
490                continue;
491            }
492            if ws > 1 && s.cpt(pos + ws) != OUT_OF_RANGE {
493                pos += ws - 1;
494                s.add(pos, &mut out);
495                continue;
496            }
497            if ws > 0 {
498                pos += ws;
499                s.add(pos, &mut out);
500                continue;
501            }
502            pos += 1;
503            s.add(pos, &mut out);
504        }
505    }
506    out
507}
508
509/// Kimi-K2: Han runs, letter words excluding Han with optional contractions, digits by three, punctuation,
510/// whitespace.
511fn kimi_k2(cpts: &[u32], offsets: &[usize]) -> Vec<usize> {
512    let mut out = Vec::with_capacity(offsets.len());
513    for mut s in spans(cpts, offsets) {
514        let letter = |s: &Span, p: usize| s.flags(p).is_letter() && !is_han(s.cpt(p));
515        let mut pos = s.lo;
516        while pos < s.hi {
517            let cpt = s.cpt(pos);
518            let flags = s.flags(pos);
519            if is_han(cpt) {
520                while is_han(s.cpt(pos)) {
521                    pos += 1;
522                }
523                s.add(pos, &mut out);
524                continue;
525            }
526            let lead_ok = !(cpt == '\r' as u32 || cpt == '\n' as u32 || flags.is_letter() || flags.is_number());
527            if (flags.is_letter() && !is_han(cpt)) || (lead_ok && letter(&s, pos + 1)) {
528                let mut has_leading = false;
529                if lead_ok {
530                    has_leading = true;
531                    pos += 1;
532                }
533                let mut has_letters = false;
534                while letter(&s, pos) {
535                    has_letters = true;
536                    pos += 1;
537                }
538                if has_letters || (!has_leading && letter(&s, pos)) {
539                    if !has_letters {
540                        pos += 1;
541                    }
542                    while letter(&s, pos) {
543                        pos += 1;
544                    }
545                    if s.cpt(pos) == '\'' as u32 && pos + 1 < s.hi {
546                        let next = to_lower(s.cpt(pos + 1));
547                        if [b's', b't', b'm', b'd'].iter().any(|&c| next == c as u32) {
548                            pos += 2;
549                        } else if pos + 2 < s.hi {
550                            let nn = to_lower(s.cpt(pos + 2));
551                            if (next == 'r' as u32 && nn == 'e' as u32)
552                                || (next == 'v' as u32 && nn == 'e' as u32)
553                                || (next == 'l' as u32 && nn == 'l' as u32)
554                            {
555                                pos += 3;
556                            }
557                        }
558                    }
559                    s.add(pos, &mut out);
560                    continue;
561                } else if has_leading {
562                    pos -= 1;
563                }
564            }
565            if flags.is_number() {
566                let mut ini = pos;
567                while s.flags(pos).is_number() {
568                    pos += 1;
569                    if pos - ini >= 3 {
570                        s.add(pos, &mut out);
571                        ini = pos;
572                    }
573                }
574                s.add(pos, &mut out);
575                continue;
576            }
577            let other = |f: Flags| !(f.is_whitespace() || f.is_letter() || f.is_number()) && f.0 != 0;
578            let mut flags2 = if cpt == ' ' as u32 { s.flags(pos + 1) } else { flags };
579            if other(flags2) {
580                pos += (cpt == ' ' as u32) as usize;
581                while other(flags2) {
582                    pos += 1;
583                    flags2 = s.flags(pos);
584                }
585                let mut c2 = s.cpt(pos);
586                while c2 == '\r' as u32 || c2 == '\n' as u32 {
587                    pos += 1;
588                    c2 = s.cpt(pos);
589                }
590                s.add(pos, &mut out);
591                continue;
592            }
593            let mut ws = 0;
594            let mut last_rn = 0;
595            while s.flags(pos + ws).is_whitespace() {
596                let c2 = s.cpt(pos + ws);
597                if c2 == '\r' as u32 || c2 == '\n' as u32 {
598                    last_rn = pos + ws + 1;
599                }
600                ws += 1;
601            }
602            if last_rn > 0 {
603                pos = last_rn;
604                s.add(pos, &mut out);
605                continue;
606            }
607            if ws > 1 && s.cpt(pos + ws) != OUT_OF_RANGE {
608                pos += ws - 1;
609                s.add(pos, &mut out);
610                continue;
611            }
612            if ws > 0 {
613                pos += ws;
614                s.add(pos, &mut out);
615                continue;
616            }
617            pos += 1;
618            s.add(pos, &mut out);
619        }
620    }
621    out
622}
623
624/// AFMoE digits: a digit run leads with its length mod 3 digits, then groups of three. The text before a run joins
625/// the run's first group, as llama.cpp's splitter emits it.
626fn afmoe(cpts: &[u32], offsets: &[usize]) -> Vec<usize> {
627    let mut out = Vec::with_capacity(offsets.len());
628    for mut s in spans(cpts, offsets) {
629        let mut pos = s.lo;
630        while pos < s.hi {
631            if s.flags(pos).is_number() {
632                let digit_start = pos;
633                let mut count = 0;
634                while s.flags(pos).is_number() && pos < s.hi {
635                    count += 1;
636                    pos += 1;
637                }
638                let rem = count % 3;
639                let mut cur = digit_start;
640                if rem > 0 {
641                    s.add(cur + rem, &mut out);
642                    cur += rem;
643                }
644                while cur < digit_start + count {
645                    s.add(cur + 3, &mut out);
646                    cur += 3;
647                }
648                continue;
649            }
650            pos += 1;
651        }
652        if s.prev_end < s.hi {
653            let hi = s.hi;
654            s.add(hi, &mut out);
655        }
656    }
657    out
658}
659
660/// `[^\n]+|[\n]+`: runs of non-newlines and runs of newlines.
661fn newlines(cpts: &[u32], offsets: &[usize]) -> Vec<usize> {
662    let mut out = Vec::with_capacity(offsets.len());
663    let mut start = 0;
664    for &offset in offsets {
665        let (lo, hi) = (start, start + offset);
666        start = hi;
667        let mut pos = lo;
668        while pos < hi {
669            let nl = cpts[pos] == '\n' as u32;
670            let run = pos;
671            while pos < hi && (cpts[pos] == '\n' as u32) == nl {
672                pos += 1;
673            }
674            out.push(pos - run);
675        }
676    }
677    out
678}
logicaffeine_compile/model/tokenizer/split.rs

logicaffeine_compile/model/tokenizer/
split.rs
