From 6f2aa0488e1b27f99c3ea9cac79e4823041f23ed Mon Sep 17 00:00:00 2001 From: AmitMY Date: Thu, 23 Jul 2026 16:34:10 +0200 Subject: [PATCH 1/2] fix: lossless merge strings via UTF-8 with byte escapes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #62. Merges crossed the string boundary through decode(errors='replace'): any non-UTF-8-aligned token (byte-level BPE constantly learns these, e.g. the first two bytes of an en-dash) collapsed to U+FFFD, so get_merges -> add_merges silently dropped it — SuperBPE's phase 2 systematically started from a weaker state than phase 1 reached. bytes_to_str/str_to_bytes make the textual form bijective: valid UTF-8 reads as itself (Hebrew, CJK, IDS composition all unchanged), literal backslashes are doubled, and undecodable bytes become \xNN escapes. GraphVertex.__str__ uses it, so display, dot labels, and get_merges agree; make_trainer inverts it exactly. Blast radius is exactly the bug: all 145 existing tests pass unchanged and 2000-row digests are identical (no invalid-byte merges at that scale); deeper training produces new (correct) digests. New shared tests: codec round-trip, and get_merges -> add_merges reproducing the exact trainer graph on a corpus with a non-UTF8-aligned merge (fails on main). Approaches rejected: latin-1 (bijective in 2 lines but cannot live in __str__ without breaking IDS lookup, and renders all non-ASCII as mojibake) and bytes-typed get_merges/add_merges (public API break, churns every string-based test). Co-Authored-By: Claude Fable 5 --- complex_tokenization/graph.py | 18 +++++++++++++++++- complex_tokenization/tokenizer.py | 4 ++-- tests/test_tokenizer_api.py | 24 ++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/complex_tokenization/graph.py b/complex_tokenization/graph.py index 77df82a..61a36cd 100644 --- a/complex_tokenization/graph.py +++ b/complex_tokenization/graph.py @@ -1,3 +1,4 @@ +import re from collections.abc import Iterable, Iterator from dataclasses import dataclass, field from functools import wraps @@ -7,6 +8,21 @@ from complex_tokenization.languages.chinese.ideographic_description_sequences import get_character_for_ids +def bytes_to_str(data: bytes) -> str: + """Lossless textual form of token bytes: valid UTF-8 reads as itself, + undecodable bytes become \\xNN escapes. Literal backslashes are doubled + first, so str_to_bytes can always invert exactly.""" + return data.replace(b"\\", b"\\\\").decode("utf-8", errors="backslashreplace") + + +def str_to_bytes(s: str) -> bytes: + return re.sub( + rb"\\\\|\\x([0-9a-fA-F]{2})", + lambda m: bytes([int(m[1], 16)]) if m[1] else b"\\", + s.encode("utf-8"), + ) + + def dot_escape(s: str) -> str: return s \ .replace("\\", "\\\\") \ @@ -44,7 +60,7 @@ def __bytes__(self): raise NotImplementedError def __str__(self): - self_str = bytes(self).decode("utf-8", errors="replace") + self_str = bytes_to_str(bytes(self)) token_replacement = get_character_for_ids(self_str) if token_replacement is not None: return token_replacement diff --git a/complex_tokenization/tokenizer.py b/complex_tokenization/tokenizer.py index c75a4a8..4c9d779 100644 --- a/complex_tokenization/tokenizer.py +++ b/complex_tokenization/tokenizer.py @@ -17,7 +17,7 @@ from tokenizers.pre_tokenizers import PreTokenizer -from complex_tokenization.graph import GraphVertex, Node +from complex_tokenization.graph import GraphVertex, Node, str_to_bytes from complex_tokenization.graphs.settings import GraphSettings from complex_tokenization.graphs.units import characters, register_script, utf8, utf8_clusters from complex_tokenization.graphs.words import GPTPretokenizer, words @@ -82,7 +82,7 @@ def make_trainer(self, texts: list[str]) -> Trainer: trainer = Trainer(graphs=graphs) for merge_strs in self.merges: - nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) + nodes = tuple(Node(value=str_to_bytes(s)) for s in merge_strs) token = reduce(lambda a, b: a + b, nodes) trainer.graph = trainer.graph.merge(token, nodes) trainer.merges.append((token, nodes)) diff --git a/tests/test_tokenizer_api.py b/tests/test_tokenizer_api.py index e4d4d8f..3ba06e7 100644 --- a/tests/test_tokenizer_api.py +++ b/tests/test_tokenizer_api.py @@ -108,3 +108,27 @@ def test_repeated_training_gives_identical_merges(self, tokenizer_cls): first = tokenizer_cls().train(texts, num_merges=4) second = tokenizer_cls().train(texts, num_merges=4) assert first == second + + +class TestMergeRoundTrip: + def test_merge_strings_reconstruct_exact_bytes(self): + from complex_tokenization.graph import bytes_to_str, str_to_bytes + for data in [b"\xe2\x80", b"the", "é".encode(), b"\\x41", b"\xff\\", " \xd7".encode("latin-1")]: + assert str_to_bytes(bytes_to_str(data)) == data + + def test_replayed_merges_reproduce_trainer_state(self): + # Byte-level BPE learns merges that are not UTF-8 aligned (here the + # first two bytes of an en-dash). get_merges -> add_merges must + # reproduce the same graph state, not silently drop those merges. + from complex_tokenization.tokenizer import BPETokenizer + texts = ["a – b – c – d – e – f – g – h"] * 3 + tok = BPETokenizer() + trainer = tok.make_trainer(texts) + tok.train_on_trainer(trainer, num_merges=6) + assert any("\\x" in s for merge in tok.get_merges() for s in merge), \ + "corpus must produce a non-UTF8-aligned merge for this test" + + tok2 = BPETokenizer() + tok2.add_merges(tok.get_merges()) + trainer2 = tok2.make_trainer(texts) + assert trainer2.graph == trainer.graph From 3fe2e3cc501b9bca3359e0513c69293af1640909 Mon Sep 17 00:00:00 2001 From: AmitMY Date: Thu, 23 Jul 2026 16:44:48 +0200 Subject: [PATCH 2/2] fix(fast): mirror the lossless merge-string codec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bytes_to_str in graph.rs is a hand-rolled UTF-8 validator walk mirroring CPython's decoder byte-for-byte (maximal-valid-subpart error ranges, per-lead-byte continuation windows, no surrogates/overlongs), with literal backslashes doubled and each invalid byte escaped as lowercase \xNN; str_to_bytes inverts with the reference's alternation order. to_str_repr uses it (IDS lookup order unchanged); the shim re-exports both from graph.py so the shared tests resolve under aliasing, and both make_trainer variants decode incoming merge strings with str_to_bytes. Fuzzed against the Python codec on 7k+ cases (random bytes, escape-heavy soup, truncated/overlong/surrogate sequences): zero mismatches, zero round-trip failures. Cross-implementation digests match under the new semantics at 2000-row/100 (unchanged values) and 20k/500 (new values, all four cases). Full-scale digests re-canonized in the PR. No perf regression (boundary-only); SuperBPE at 5000 merges is 1.3s faster — with no merges silently dropped, phase 2 starts from the true, more-merged phase-1 state and does less work. Co-Authored-By: Claude Fable 5 --- .../python/complex_tokenization_fast/graph.py | 4 +- .../complex_tokenization_fast/tokenizer.py | 6 +- fast/src/graph.rs | 102 +++++++++++++++++- fast/src/lib.rs | 2 + 4 files changed, 108 insertions(+), 6 deletions(-) diff --git a/fast/python/complex_tokenization_fast/graph.py b/fast/python/complex_tokenization_fast/graph.py index 896d23f..9708100 100644 --- a/fast/python/complex_tokenization_fast/graph.py +++ b/fast/python/complex_tokenization_fast/graph.py @@ -1,9 +1,11 @@ -from complex_tokenization_fast._rs import ( +from complex_tokenization_fast._rs import ( # noqa: F401 FullyConnectedGraph, Node, NodesSequence, Tree, UnconnectedGraphs, + bytes_to_str, + str_to_bytes, ) GraphVertex = (Node, NodesSequence, Tree, FullyConnectedGraph, UnconnectedGraphs) diff --git a/fast/python/complex_tokenization_fast/tokenizer.py b/fast/python/complex_tokenization_fast/tokenizer.py index d84b2bf..9653ff2 100644 --- a/fast/python/complex_tokenization_fast/tokenizer.py +++ b/fast/python/complex_tokenization_fast/tokenizer.py @@ -1,6 +1,6 @@ from functools import lru_cache, reduce -from complex_tokenization_fast._rs import Node, Trainer +from complex_tokenization_fast._rs import Node, Trainer, str_to_bytes from complex_tokenization_fast.graphs.settings import GraphSettings from complex_tokenization_fast.graphs.units import characters, register_script, utf8, utf8_clusters from complex_tokenization_fast.graphs.words import GPTPretokenizer, words @@ -56,7 +56,7 @@ def make_trainer(self, texts): if self.merges: merge_list = [] for merge_strs in self.merges: - nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) + nodes = tuple(Node(value=str_to_bytes(s)) for s in merge_strs) token = reduce(lambda a, b: a + b, nodes) merge_list.append((token, nodes)) trainer.apply_merges(merge_list) @@ -81,7 +81,7 @@ def make_streaming_trainer(self, texts): if self.merges: merge_list = [] for merge_strs in self.merges: - nodes = tuple(Node(value=s.encode("utf-8")) for s in merge_strs) + nodes = tuple(Node(value=str_to_bytes(s)) for s in merge_strs) token = reduce(lambda a, b: a + b, nodes) merge_list.append((token, nodes)) trainer.apply_merges(merge_list) diff --git a/fast/src/graph.rs b/fast/src/graph.rs index ad408f3..06cac05 100644 --- a/fast/src/graph.rs +++ b/fast/src/graph.rs @@ -82,6 +82,93 @@ impl Hash for GraphV { } } +/// Lossless textual form of token bytes, byte-for-byte equal to the +/// reference's `bytes_to_str`: literal backslashes doubled, valid UTF-8 runs +/// kept as-is, and each byte of an invalid sequence escaped as lowercase +/// \xNN. Error ranges mirror CPython's UTF-8 decoder (maximal valid +/// subpart): the start byte plus any valid continuations seen before the +/// failure form one range, escaped per byte. +pub fn bytes_to_str(data: &[u8]) -> String { + // Length of the valid UTF-8 sequence at data[i..], or the error-range + // length as Err (>= 1). Continuation windows per CPython: E0 A0-BF, + // ED 80-9F, F0 90-BF, F4 80-8F, plain 80-BF elsewhere. + fn seq_len(data: &[u8], i: usize) -> Result { + let b = data[i]; + let (n, first_lo, first_hi) = match b { + 0x00..=0x7f => return Ok(1), + 0xc2..=0xdf => (2, 0x80, 0xbf), + 0xe0 => (3, 0xa0, 0xbf), + 0xe1..=0xec | 0xee..=0xef => (3, 0x80, 0xbf), + 0xed => (3, 0x80, 0x9f), + 0xf0 => (4, 0x90, 0xbf), + 0xf1..=0xf3 => (4, 0x80, 0xbf), + 0xf4 => (4, 0x80, 0x8f), + _ => return Err(1), + }; + for k in 1..n { + let (lo, hi) = if k == 1 { (first_lo, first_hi) } else { (0x80, 0xbf) }; + match data.get(i + k) { + Some(&c) if c >= lo && c <= hi => {} + _ => return Err(k), + } + } + Ok(n) + } + + let mut out = String::with_capacity(data.len() + 8); + let mut i = 0; + while i < data.len() { + if data[i] == b'\\' { + out.push_str("\\\\"); + i += 1; + continue; + } + match seq_len(data, i) { + Ok(n) => { + out.push_str(std::str::from_utf8(&data[i..i + n]).unwrap()); + i += n; + } + Err(n) => { + for &b in &data[i..i + n] { + out.push_str(&format!("\\x{b:02x}")); + } + i += n; + } + } + } + out +} + +/// Exact inverse of `bytes_to_str`: on the UTF-8 bytes of `s`, a doubled +/// backslash becomes one backslash, \xNN becomes that byte (doubled +/// backslash wins, matching the reference's regex alternation order). +pub fn str_to_bytes(s: &str) -> Vec { + let b = s.as_bytes(); + let mut out = Vec::with_capacity(b.len()); + let mut i = 0; + while i < b.len() { + if b[i] == b'\\' { + if b.get(i + 1) == Some(&b'\\') { + out.push(b'\\'); + i += 2; + continue; + } + if b.get(i + 1) == Some(&b'x') + && b.get(i + 2).is_some_and(u8::is_ascii_hexdigit) + && b.get(i + 3).is_some_and(u8::is_ascii_hexdigit) + { + let hex = std::str::from_utf8(&b[i + 2..i + 4]).unwrap(); + out.push(u8::from_str_radix(hex, 16).unwrap()); + i += 4; + continue; + } + } + out.push(b[i]); + i += 1; + } + out +} + impl GraphV { pub fn to_bytes(&self) -> Vec { match self { @@ -114,8 +201,7 @@ impl GraphV { } pub fn to_str_repr(&self) -> String { - let bytes = self.to_bytes(); - let s = String::from_utf8_lossy(&bytes).to_string(); + let s = bytes_to_str(&self.to_bytes()); // Try IDS character lookup if let Some(ch) = crate::units::get_character_for_ids_str(&s) { return ch; @@ -563,6 +649,18 @@ fn fullconn_merge(nodes: &[GraphV], token: &GraphV, merge: &[GraphV]) -> GraphV GraphV::new_fullconn(new_nodes) } +#[pyfunction] +#[pyo3(name = "bytes_to_str")] +pub fn bytes_to_str_py(data: Vec) -> String { + bytes_to_str(&data) +} + +#[pyfunction] +#[pyo3(name = "str_to_bytes")] +pub fn str_to_bytes_py(py: Python<'_>, s: &str) -> PyObject { + pyo3::types::PyBytes::new(py, &str_to_bytes(s)).into() +} + // --- Python wrapper types --- pub fn graphv_to_pyobject(py: Python<'_>, g: &GraphV) -> PyObject { diff --git a/fast/src/lib.rs b/fast/src/lib.rs index 842a1b9..c025895 100644 --- a/fast/src/lib.rs +++ b/fast/src/lib.rs @@ -22,6 +22,8 @@ fn _rs(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(units::set_ids_reverse_dict_py, m)?)?; m.add_function(wrap_pyfunction!(units::warm_word_cache_py, m)?)?; m.add_function(wrap_pyfunction!(units::clear_word_cache, m)?)?; + m.add_function(wrap_pyfunction!(graph::bytes_to_str_py, m)?)?; + m.add_function(wrap_pyfunction!(graph::str_to_bytes_py, m)?)?; m.add_function(wrap_pyfunction!(sync_settings, m)?)?; Ok(()) }