fix(fast): training no longer mutates the global cluster cache

CLAUDE.md

@@ -18,11 +18,11 @@ To verify:
 pip install -e . && pytest tests/
 
 # Test fast (from fast/ directory)
-cd fast && maturin develop --release && pytest ../tests/
+cd fast && maturin develop --release && pytest ../tests/ -p conftest
 ```
 
 ## Testing
 
 - Reference tests run from the repo root: `pytest tests/`
-- Fast tests run from `fast/`: `pytest ../tests/` (the `fast/conftest.py` aliases `complex_tokenization` → `complex_tokenization_fast`)
+- Fast tests run from `fast/`: `pytest ../tests/ -p conftest` — the `-p conftest` is REQUIRED: it loads `fast/conftest.py`, which aliases `complex_tokenization` → `complex_tokenization_fast`. Without it, pytest 9 never loads that conftest (it is outside the tests/ subtree) and the run silently tests the reference implementation instead of fast
 - Both must pass the same 127+ tests

fast/python/complex_tokenization_fast/graphs/settings.py

@@ -4,6 +4,9 @@
 class _GraphSettingsMeta(type):
     _MAX_MERGE_SIZE = 2
     _ONLY_MINIMAL_MERGES = True
+    # Reference-implementation knob (gates its get_merges memoization); the
+    # fast implementation has no equivalent cache, so it only round-trips.
+    TRADE_MEMORY_FOR_SPEED = True
 
     @property
     def MAX_MERGE_SIZE(cls):  # noqa: N802

fast/python/complex_tokenization_fast/tokenizer.py

@@ -1,4 +1,4 @@
-from functools import reduce
+from functools import lru_cache, reduce
 
 from complex_tokenization_fast._rs import Node, Trainer
 from complex_tokenization_fast.graphs.settings import GraphSettings
@@ -13,7 +13,8 @@
 
 
 class Tokenizer:
-    def __init__(self, units="utf8_clusters", merge_size=2, connected=False, pretokenizer=GPTPretokenizer):
+    def __init__(self, units="utf8_clusters", merge_size=2, connected=False, pretokenizer=GPTPretokenizer,
+                 cache_maxsize=None):
         if isinstance(units, str):
             if units not in UNIT_FUNCTIONS:
                 raise ValueError(f"Unknown units: {units!r}. Choose from {list(UNIT_FUNCTIONS)}")
@@ -23,6 +24,7 @@ def __init__(self, units="utf8_clusters", merge_size=2, connected=False, pretoke
         self.merge_size = merge_size
         self.connected = connected
         self.pretokenizer = pretokenizer
+        self.cache_maxsize = cache_maxsize
         self.merges = []
 
     @staticmethod
@@ -33,8 +35,14 @@ def add_merges(self, merges):
         self.merges.extend(merges)
 
     def _build_graphs(self, texts):
+        # Same build-local word-graph dedup as the reference: repeated words
+        # share one graph. cache_maxsize=None is unbounded; 0 disables.
+        if self.cache_maxsize == 0:
+            units = self.units
+        else:
+            units = lru_cache(maxsize=self.cache_maxsize)(self.units)
         return tuple(
-            words(text, connected=self.connected, units=self.units, pretokenizer=self.pretokenizer)
+            words(text, connected=self.connected, units=units, pretokenizer=self.pretokenizer)
             for text in texts
         )
 
@@ -92,28 +100,28 @@ def get_merges(self):
 
 
 class BPETokenizer(Tokenizer):
-    def __init__(self, units="utf8_clusters", pretokenizer=GPTPretokenizer):
-        super().__init__(units=units, merge_size=2, connected=False, pretokenizer=pretokenizer)
+    def __init__(self, **kwargs):
+        super().__init__(merge_size=2, connected=False, **kwargs)
 
 
 class BNETokenizer(Tokenizer):
-    def __init__(self, n=4, units="utf8_clusters", pretokenizer=GPTPretokenizer):
-        super().__init__(units=units, merge_size=n, connected=False, pretokenizer=pretokenizer)
+    def __init__(self, n=4, **kwargs):
+        super().__init__(merge_size=n, connected=False, **kwargs)
 
 
 class BoundlessBPETokenizer(Tokenizer):
-    def __init__(self, units="utf8_clusters", pretokenizer=GPTPretokenizer):
-        super().__init__(units=units, merge_size=2, connected=True, pretokenizer=pretokenizer)
+    def __init__(self, **kwargs):
+        super().__init__(merge_size=2, connected=True, **kwargs)
 
 
 class SuperBPETokenizer(Tokenizer):
-    def __init__(self, units="utf8_clusters", disconnected_merges=None, pretokenizer=GPTPretokenizer):
-        super().__init__(units=units, merge_size=2, connected=False, pretokenizer=pretokenizer)
+    def __init__(self, disconnected_merges=None, **kwargs):
+        super().__init__(merge_size=2, connected=False, **kwargs)
         self._disconnected_merges = disconnected_merges
 
     def train(self, texts, num_merges=100, progress=False):
         disconnected_merges = self._disconnected_merges or num_merges // 2
-        phase1 = BPETokenizer(units=self.units, pretokenizer=self.pretokenizer)
+        phase1 = BPETokenizer(units=self.units, pretokenizer=self.pretokenizer, cache_maxsize=self.cache_maxsize)
         phase1.train(texts, num_merges=disconnected_merges)
         self.connected = True
         self.add_merges(phase1.merges)

fast/src/trainer.rs

@@ -5,7 +5,7 @@ use rustc_hash::{FxHashMap, FxHashSet};
 use std::sync::Arc;
 
 use crate::graph::{graphv_to_pyobject, pyobject_to_graphv, GraphV};
-use crate::units::{apply_merge_to_cluster_cache, replace_word_cache, snapshot_word_cache};
+use crate::units::{replace_word_cache, snapshot_word_cache};
 use std::collections::HashMap;
 
 // Merge score: merging a k-tuple with `count` occurrences removes (k-1)*count nodes.
@@ -430,9 +430,6 @@ fn train_unconn(
         num_merges,
         verbose,
         true,
-        &mut |token, nodes| {
-            apply_merge_to_cluster_cache(token, nodes);
-        },
     );
 
     for (entry, member_idxs) in entries.iter().zip(&members) {
@@ -462,7 +459,6 @@ fn train_single(
             println!("Merging {:?} count={}", nodes, counts[&nodes]);
         }
         *graph = graph.merge(&token, &nodes);
-        apply_merge_to_cluster_cache(&token, &nodes);
         pending.push((token, nodes));
     }
     pending
@@ -846,7 +842,6 @@ fn build_candidate_index(entries: &[WordEntry]) -> Vec<FxHashMap<Vec<GraphV>, Fx
 // Persistent global counts + inverted index with delta updates. Per merge step
 // only the entries containing the chosen candidate are touched, and `global` is
 // patched by diffing each entry's old vs new candidates instead of being rebuilt.
-// `on_merge` fires once per merge in order (used for cluster-cache side effects).
 // `use_try` selects try_merge semantics (legacy Unconn path) over merge; with
 // try_merge a no-op leaves the entry and counts untouched.
 fn train_entries_delta(
@@ -856,7 +851,6 @@ fn train_entries_delta(
     num_merges: usize,
     verbose: bool,
     use_try: bool,
-    on_merge: &mut dyn FnMut(&GraphV, &[GraphV]),
 ) -> Vec<(GraphV, Vec<GraphV>)> {
     let mut pending = Vec::new();
     let mut global = build_global_counts(entries);
@@ -1043,7 +1037,6 @@ fn train_entries_delta(
                     }
                 }
             });
-        on_merge(&token, &nodes);
         pending.push((token, nodes));
     }
     pending
@@ -1056,7 +1049,7 @@ fn train_streaming_disconnected(
     num_merges: usize,
     verbose: bool,
 ) -> Vec<(GraphV, Vec<GraphV>)> {
-    train_entries_delta(entries, None, range_start, num_merges, verbose, false, &mut |_, _| {})
+    train_entries_delta(entries, None, range_start, num_merges, verbose, false)
 }
 
 // Connected: must assemble doc graphs for cross-word pairs
@@ -1406,7 +1399,6 @@ impl Trainer {
                     };
                     let token = make_token(&nodes);
                     apply_merge_parallel(subs, &mut active, &token, &nodes);
-                    apply_merge_to_cluster_cache(&token, &nodes);
                     pending.push((token, nodes));
 
                     let merge_num = i + 1;
@@ -1426,7 +1418,6 @@ impl Trainer {
                     };
                     let token = make_token(&nodes);
                     *graph = graph.merge(&token, &nodes);
-                    apply_merge_to_cluster_cache(&token, &nodes);
                     pending.push((token, nodes));
 
                     let merge_num = i + 1;

fast/src/units.rs

@@ -52,15 +52,6 @@ pub fn clear_handlers() {
     CLUSTER_CACHE.lock().unwrap().clear();
 }
 
-pub fn apply_merge_to_cluster_cache(token: &GraphV, merge: &[GraphV]) {
-    let mut cache = CLUSTER_CACHE.lock().unwrap();
-    for graph in cache.values_mut() {
-        if let Some(new_g) = graph.try_merge(token, merge) {
-            *graph = new_g;
-        }
-    }
-}
-
 static WORD_CACHE: LazyLock<Mutex<HashMap<String, GraphV>>> =
     LazyLock::new(|| Mutex::new(HashMap::new()));
 

tests/test_tokenizer_api.py

@@ -91,3 +91,20 @@ def test_different_pretokenizer_different_merges(self):
         default_merges = BPETokenizer().train(texts, num_merges=5)
         whitespace_merges = BPETokenizer(pretokenizer=Whitespace()).train(texts, num_merges=5)
         assert default_merges != whitespace_merges
+
+    @pytest.mark.parametrize("tokenizer_cls", [
+        BPETokenizer,
+        lambda: BNETokenizer(n=4),
+        BoundlessBPETokenizer,
+        SuperBPETokenizer,
+    ])
+    def test_repeated_training_gives_identical_merges(self, tokenizer_cls):
+        # Training must not observably mutate process-global state: a second
+        # train in the same process must reproduce the first. The corpus has
+        # multi-byte characters (é, à) whose UTF-8 byte pairs get merged
+        # within the first few merges — a cached cluster graph mutated by the
+        # first train would change what the second train sees.
+        texts = ["café déjà café déjà café déjà"]
+        first = tokenizer_cls().train(texts, num_merges=4)
+        second = tokenizer_cls().train(texts, num_merges=4)
+        assert first == second