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/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(()) } 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