Simpler and faster encoding
Author: milan-kodaliCreated Oct 1, 2025Updated Nov 9, 2025
I think the logic in encode can be simplified, making it easier to understand, while meaningfully improving performance.
The current function, in my understanding, should have runtime proportional to 3*m*n where m is the length of merges and n is the length of text.
def encode(self, text):
text_bytes = text.encode("utf-8")
ids = list(text_bytes)
while len(ids) >= 2: # should loop ~m times
stats = get_stats(ids) # O(n)
pair = min(stats, key=lambda p: self.merges.get(p, float("inf"))) # O(n)
if pair not in self.merges:
break
idx = self.merges[pair]
ids = merge(ids, pair, idx) # O(n)
return idsWe can simplify the logic significantly by leveraging the existing merge function and just merging over every pair in merges. This should also have a runtime proportional to 1*m*n:
def faster_encode(self, text):
text_bytes = text.encode("utf-8")
ids = list(text_bytes)
for pair, idx in self.merges.items(): # O(m)
ids = merge(ids, pair, idx) # O(n)
return idsIn testing with various vocab sizes and input text length, I'm seeing a ~2-3x speed improvement for faster_encode, and it seems like this logic can be extended to the other Tokenizer classes.
Source: karpathy/minbpe