use super::unicode_class::{Tables, tables};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ScanKind {
O200k,
Kimi,
}
pub(crate) fn recognize(source: &str) -> Option<ScanKind> {
if source == crate::tiktoken::O200K_BASE_PATTERN {
Some(ScanKind::O200k)
} else if source == crate::tiktoken::KIMI_PATTERN {
Some(ScanKind::Kimi)
} else {
None
}
}
#[inline(always)]
fn is_letter(t: &Tables, c: char) -> bool {
let u = c as u32;
if u < 0x80 {
c.is_ascii_alphabetic()
} else {
t.is_letter(u)
}
}
#[inline(always)]
fn is_number(t: &Tables, c: char) -> bool {
let u = c as u32;
if u < 0x80 {
c.is_ascii_digit()
} else {
t.is_number(u)
}
}
#[inline(always)]
fn is_ws(t: &Tables, c: char) -> bool {
let u = c as u32;
if u < 0x80 {
matches!(u as u8, b'\t' | b'\n' | 0x0b | 0x0c | b'\r' | b' ')
} else {
t.is_ws(u)
}
}
#[inline(always)]
fn is_han(t: &Tables, c: char) -> bool {
let u = c as u32;
if u < 0x80 { false } else { t.is_han(u) }
}
#[inline(always)]
fn is_ugroup(t: &Tables, c: char, kimi: bool) -> bool {
let u = c as u32;
let m = if u < 0x80 {
c.is_ascii_uppercase()
} else {
t.is_ugroup(u)
};
m && !(kimi && is_han(t, c))
}
#[inline(always)]
fn is_lgroup(t: &Tables, c: char, kimi: bool) -> bool {
let u = c as u32;
let m = if u < 0x80 {
c.is_ascii_lowercase()
} else {
t.is_lgroup(u)
};
m && !(kimi && is_han(t, c))
}
#[inline(always)]
fn run_member(t: &Tables, c: char, kimi: bool) -> bool {
is_ugroup(t, c, kimi) || is_lgroup(t, c, kimi)
}
#[inline(always)]
fn is_prefix(t: &Tables, c: char) -> bool {
c != '\r' && c != '\n' && !is_letter(t, c) && !is_number(t, c)
}
#[inline(always)]
fn is_punct(t: &Tables, c: char) -> bool {
!is_ws(t, c) && !is_letter(t, c) && !is_number(t, c)
}
#[inline]
fn contraction_len(bs: &[u8]) -> usize {
if bs.len() < 2 || bs[1] >= 0x80 {
return 0;
}
match bs[1] | 0x20 {
b's' | b't' | b'm' | b'd' => 2,
b'r' | b'v' if bs.len() >= 3 && (bs[2] | 0x20) == b'e' => 3,
b'l' if bs.len() >= 3 && (bs[2] | 0x20) == b'l' => 3,
_ => 0,
}
}
#[inline(always)]
fn char_at(text: &str, b: &[u8], i: usize) -> (char, usize) {
let byte = b[i];
if byte < 0x80 {
(byte as char, 1)
} else {
let c = text[i..].chars().next().unwrap();
(c, c.len_utf8())
}
}
const SWAR_HI: u64 = 0x8080_8080_8080_8080;
#[inline(always)]
fn ascii_lower_run_end(b: &[u8], mut pos: usize) -> usize {
let n = b.len();
while pos + 8 <= n {
let word = unsafe { (b.as_ptr().add(pos) as *const u64).read_unaligned() };
if word & SWAR_HI != 0 {
break; }
let ge_a = (word | SWAR_HI).wrapping_sub(0x6161_6161_6161_6161);
let le_z = 0xFAFA_FAFA_FAFA_FAFA_u64.wrapping_sub(word);
let non_lower = !(ge_a & le_z) & SWAR_HI;
if non_lower != 0 {
return pos + non_lower.to_le().trailing_zeros() as usize / 8;
}
pos += 8;
}
while pos < n {
let x = unsafe { *b.get_unchecked(pos) };
if x.wrapping_sub(b'a') < 26 {
pos += 1;
} else {
break;
}
}
pos
}
pub(crate) fn newline_chunk_bounds(text: &str, n_chunks: usize) -> Vec<(usize, usize)> {
let bytes = text.as_bytes();
let n = bytes.len();
if n_chunks < 2 {
return vec![(0, n)];
}
let t = tables();
let nominal = n / n_chunks;
let mut splits = vec![0usize];
for i in 1..n_chunks {
let from = i * nominal;
let Some(rel) = memchr::memchr2(b'\n', b'\r', &bytes[from..]) else {
break;
};
let mut e = from + rel;
let mut last_nl = e;
while e < n {
let (c, l) = char_at(text, bytes, e);
if !is_ws(t, c) {
break;
}
if bytes[e] == b'\n' || bytes[e] == b'\r' {
last_nl = e;
}
e += l;
}
let boundary = last_nl + 1;
if boundary < n && boundary > *splits.last().unwrap() {
splits.push(boundary);
}
}
splits.push(n);
splits.windows(2).map(|w| (w[0], w[1])).collect()
}
pub(crate) fn scan_core<F>(kind: ScanKind, text: &str, mut emit: F) -> Result<(), String>
where
F: FnMut(usize, usize) -> Result<(), String>,
{
let t = tables();
let b = text.as_bytes();
let n = b.len();
let kimi = kind == ScanKind::Kimi;
let slash = kind == ScanKind::O200k;
let mut i = 0usize;
while i < n {
let start = i;
{
let c0 = b[i];
let lstart = if c0.wrapping_sub(b'a') < 26 {
i
} else if c0 == b' ' && i + 1 < n && b[i + 1].wrapping_sub(b'a') < 26 {
i + 1
} else {
usize::MAX
};
if lstart != usize::MAX {
let run_end = ascii_lower_run_end(b, lstart);
if run_end == n || (b[run_end] < 0x80 && b[run_end] != b'\'') {
emit(start, run_end)?;
i = run_end;
continue;
}
}
}
let (c, clen) = char_at(text, b, i);
if kimi && is_han(t, c) {
let mut e = i + clen;
while e < n {
let (cj, lj) = char_at(text, b, e);
if is_han(t, cj) {
e += lj;
} else {
break;
}
}
emit(start, e)?;
i = e;
continue;
}
let run_start = if run_member(t, c, kimi) {
Some(i)
} else if is_prefix(t, c) && i + clen < n {
let (c1, _) = char_at(text, b, i + clen);
if run_member(t, c1, kimi) {
Some(i + clen)
} else {
None
}
} else {
None
};
if let Some(run_start) = run_start {
let mut u = run_start;
let mut last_lg = usize::MAX;
while u < n {
let (cj, lj) = char_at(text, b, u);
if !is_ugroup(t, cj, kimi) {
break;
}
if is_lgroup(t, cj, kimi) {
last_lg = u;
}
u += lj;
}
let after_u_is_lgroup = u < n && is_lgroup(t, char_at(text, b, u).0, kimi);
let lstart = if after_u_is_lgroup { u } else { last_lg };
let mut e = if lstart != usize::MAX {
let mut e = lstart;
while e < n {
let (cj, lj) = char_at(text, b, e);
if !is_lgroup(t, cj, kimi) {
break;
}
e += lj;
}
e
} else {
u
};
if e < n && b[e] == b'\'' {
e += contraction_len(&b[e..]);
}
emit(start, e)?;
i = e;
continue;
}
if is_number(t, c) {
let mut cnt = 1usize;
let mut e = i + clen;
while cnt < 3 && e < n {
let (cj, lj) = char_at(text, b, e);
if is_number(t, cj) {
cnt += 1;
e += lj;
} else {
break;
}
}
emit(start, e)?;
i = e;
continue;
}
{
let (mut pstart, mut pc, mut pcl) = (i, c, clen);
if c == ' ' && i + clen < n {
let (c1, l1) = char_at(text, b, i + clen);
if is_punct(t, c1) {
pstart = i + clen;
pc = c1;
pcl = l1;
}
}
if is_punct(t, pc) {
let mut e = pstart + pcl;
while e < n {
let (cj, lj) = char_at(text, b, e);
if is_punct(t, cj) {
e += lj;
} else {
break;
}
}
while e < n && (b[e] == b'\r' || b[e] == b'\n' || (slash && b[e] == b'/')) {
e += 1;
}
emit(start, e)?;
i = e;
continue;
}
}
let mut e = i;
let mut last_cp_start = i;
while e < n {
let (cj, lj) = char_at(text, b, e);
if is_ws(t, cj) {
last_cp_start = e;
e += lj;
} else {
break;
}
}
let we = e;
let mut last_nl = usize::MAX;
let mut k = i;
while k < we {
if b[k] == b'\r' || b[k] == b'\n' {
last_nl = k;
}
k += 1;
}
let end = if last_nl != usize::MAX {
last_nl + 1 } else if we == n {
we } else if last_cp_start > i {
last_cp_start } else {
we };
emit(i, end)?;
i = end;
}
Ok(())
}
#[cfg(test)]
fn scan_seq(kind: ScanKind, text: &str) -> Vec<(u32, u32)> {
let mut out: Vec<(u32, u32)> = Vec::with_capacity(text.len() / 4 + 1);
let _ = scan_core(kind, text, |s, e| {
out.push((s as u32, e as u32));
Ok(())
});
out
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tiktoken::KIMI_PATTERN;
fn scan(kind: ScanKind, text: &str) -> Vec<String> {
scan_seq(kind, text)
.iter()
.map(|&(s, e)| text[s as usize..e as usize].to_string())
.collect()
}
#[test]
fn chunk_bounds_preserve_interior_newline_pretokens() {
let text = format!("{} \n \n{}", "a".repeat(100), "b".repeat(95));
let whole = scan(ScanKind::Kimi, &text);
assert!(
whole.iter().any(|p| p == " \n \n"),
"run should be one pretoken: {whole:?}"
);
let bounds = newline_chunk_bounds(&text, 2);
assert!(bounds.len() >= 2, "expected a split: {bounds:?}");
let chunked: Vec<String> = bounds
.iter()
.flat_map(|&(s, e)| scan(ScanKind::Kimi, &text[s..e]))
.collect();
assert_eq!(chunked, whole, "chunked scan diverged from whole scan");
}
#[test]
fn words_case_split() {
assert_eq!(scan(ScanKind::O200k, "HTTPRequest"), vec!["HTTPRequest"]);
assert_eq!(scan(ScanKind::O200k, "HelloWorld"), vec!["Hello", "World"]);
assert_eq!(scan(ScanKind::O200k, "camelCase"), vec!["camel", "Case"]);
assert_eq!(scan(ScanKind::O200k, "iOS"), vec!["i", "OS"]);
assert_eq!(scan(ScanKind::O200k, "ALLCAPS"), vec!["ALLCAPS"]);
assert_eq!(scan(ScanKind::O200k, "aBc"), vec!["a", "Bc"]);
}
#[test]
fn contractions_and_prefix() {
assert_eq!(scan(ScanKind::O200k, "don't"), vec!["don't"]);
assert_eq!(scan(ScanKind::O200k, "I'll"), vec!["I'll"]);
assert_eq!(scan(ScanKind::O200k, "O'Brien"), vec!["O", "'Brien"]);
assert_eq!(
scan(ScanKind::O200k, "hello world"),
vec!["hello", " world"]
);
assert_eq!(scan(ScanKind::O200k, "a!b"), vec!["a", "!b"]);
}
#[test]
fn numbers_punct_whitespace() {
assert_eq!(scan(ScanKind::O200k, "1234567"), vec!["123", "456", "7"]);
assert_eq!(scan(ScanKind::O200k, "3.14"), vec!["3", ".", "14"]);
assert_eq!(scan(ScanKind::O200k, "!!!"), vec!["!!!"]);
assert_eq!(scan(ScanKind::O200k, "a b"), vec!["a", " ", " b"]);
assert_eq!(scan(ScanKind::O200k, "trailing "), vec!["trailing", " "]);
assert_eq!(
scan(ScanKind::O200k, "foo\n\nbar"),
vec!["foo", "\n\n", "bar"]
);
assert_eq!(scan(ScanKind::O200k, " \n x"), vec![" \n", " ", " x"]);
}
#[test]
fn unicode_and_han() {
assert_eq!(scan(ScanKind::Kimi, "café"), vec!["café"]);
assert_eq!(scan(ScanKind::Kimi, "你好world"), vec!["你好", "world"]);
assert_eq!(
scan(ScanKind::Kimi, "café中文test"),
vec!["café", "中文", "test"]
);
assert_eq!(scan(ScanKind::Kimi, "中1文"), vec!["中", "1", "文"]);
assert_eq!(scan(ScanKind::Kimi, "a\u{3000}b"), vec!["a", "\u{3000}b"]);
assert_eq!(scan(ScanKind::O200k, "你好"), vec!["你好"]);
}
#[test]
fn newline_chunking_matches_whole() {
let unit = "Hello world!\nCamelCase 中文 test\n\n spaced lines \n \n\
café résumé 12345 don't \n \t\n更多文本\r\n";
let big = unit.repeat(400);
for kind in [ScanKind::O200k, ScanKind::Kimi] {
let whole = scan_seq(kind, &big);
for n_chunks in [1usize, 2, 3, 7, 16, 64] {
let mut combined = Vec::new();
for (s, e) in newline_chunk_bounds(&big, n_chunks) {
let base = s as u32;
for (a, b) in scan_seq(kind, &big[s..e]) {
combined.push((a + base, b + base));
}
}
assert_eq!(combined, whole, "kind={kind:?} n_chunks={n_chunks}");
}
}
}
#[test]
fn recognizes_patterns() {
assert_eq!(
recognize(crate::tiktoken::O200K_BASE_PATTERN),
Some(ScanKind::O200k)
);
assert_eq!(recognize(KIMI_PATTERN), Some(ScanKind::Kimi));
assert_eq!(recognize("something else"), None);
}
#[test]
fn scanner_matches_regex_engine() {
use crate::Split;
use crate::pre_tokenized::PreTokenizedString;
use serde_json::json;
let corpus = [
"",
"Hello, world! HTTPRequest HelloWorld camelCase ALLCAPS iOS getHTTPResponse",
"don't I'll O'Brien y'all wasn't 'tis can't won't",
"1234567 3.14 1,000 42 007 mixed ABC123def a1b2c3",
" leading trailing a b a b\ttabs\there",
"foo\n\nbar \n x\r\n\r\nwin end. x.\n\n",
"!!! ... @#$ a!b ( hello .. word e.g. U.S.A. snake_case kebab-case -42 ' -42",
"café résumé naïve über straße señor niño",
"你好世界 中文,世界 混合ABCと日本語 test中文HELLO café中文test 中1文 a中b",
"こんにちは 안녕하세요 Привет мир Ελληνικά עברית العربية",
"數據科學 fêteliefer.githubusercontent શહેરefeller эндey feature",
"עצמ亚洲AVurant בדרך无码AVJobҮ下面tellremaining მჯდომ",
"😀🚀 中文 test 한국어 test 中文 日本語 ①②③ Full",
"a\u{3000}b \u{00a0}nbsp \u{2028}line é\u{0301}\u{0302} Dž titlecase",
"The rain in Spain. ".to_string().repeat(50).leak(),
];
for (kind, pat) in [
(ScanKind::O200k, crate::tiktoken::O200K_BASE_PATTERN),
(ScanKind::Kimi, KIMI_PATTERN),
] {
let split = Split::from_config(&json!({ "Regex": pat }), "Isolated", false).unwrap();
for s in &corpus {
let mut pts = PreTokenizedString::from_text(s);
split.pre_tokenize(&mut pts).unwrap();
let regex_ranges: Vec<(u32, u32)> = pts
.splits()
.iter()
.filter(|sp| !sp.range.is_empty())
.map(|sp| (sp.range.start as u32, sp.range.end as u32))
.collect();
let scan: Vec<(u32, u32)> = scan_seq(kind, s)
.into_iter()
.filter(|(a, b)| a != b)
.collect();
assert_eq!(scan, regex_ranges, "kind={kind:?} input={s:?}");
}
}
}
}