fix: lossless merge strings via UTF-8 with byte escapes

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

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))

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)

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)

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<usize, usize> {
+        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<u8> {
+    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<u8> {
         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<u8>) -> 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 {

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

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