Skip to content

Commit fb7b73f

Browse files
AmitMYclaude
andcommitted
perf: merge and recount each unique subgraph object once in incremental loop
Duplicated words share one subgraph object via the build cache, but the incremental update loop merged and recounted per occurrence, so duplicates diverged into equal-but-distinct copies on their first merge. Memoizing by object identity does the work once per unique word; the resulting Counter is shared across occurrences (comp_counts entries are only ever read). Zipf corpus (20k words, 2k vocab, 300 merges), byte-identical output: - BPE: 0.362s -> 0.172s, peak 35.8MB -> 31.5MB - BNE: 0.501s -> 0.352s, peak 37.8MB -> 35.3MB Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 25b3773 commit fb7b73f

1 file changed

Lines changed: 9 additions & 2 deletions

File tree

‎complex_tokenization/trainer.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,10 +95,17 @@ def _train_incremental(self, num_merges: int, progress: bool = False):
9595
nodes = next(m for cc in comp_counts for m in cc if (len(m) - 1) * total[m] == best)
9696
token = reduce(lambda x, y: x + y, nodes)
9797

98+
# Duplicated words share one subgraph object (the build cache), so
99+
# merge and recount each distinct object once, not per occurrence.
100+
# The shared Counter is safe: comp_counts entries are only read.
101+
merged: dict[int, tuple[GraphVertex, Counter]] = {}
98102
for i in list(index[nodes]):
99103
_index_remove(total, index, i, comp_counts[i])
100-
components[i] = components[i].merge(token, nodes)
101-
comp_counts[i] = Counter(components[i].get_merges())
104+
pair = merged.get(id(components[i]))
105+
if pair is None:
106+
component = components[i].merge(token, nodes)
107+
pair = merged[id(components[i])] = (component, Counter(component.get_merges()))
108+
components[i], comp_counts[i] = pair
102109
_index_add(total, index, i, comp_counts[i])
103110

104111
self.merges.append((token, nodes))

0 commit comments

Comments
 (0)