The whole training run with pair counts and ties, then the learned rules applied to words the corpus never held.
The notebook carries the source, the install step and every run below, so nothing has to be on your machine. Open it, then choose Runtime › Run all.
Colab opens it read only. Click Copy to Drive to keep
your changes. You can also download
bpe.ipynb, or run it
locally: python3 bpe.py.
python3 bpe.py
The five merges, each with the counts that decided it.
python3 bpe.py --no-tiebreak
Take the lexicographically last tied pair instead. Three of the five steps involve a tie, so every merge changes and so does every tokenisation. The rule is not a footnote.
python3 bpe.py --encode lowest
Two clean pieces, because lowest is built from parts the corpus paid for.
python3 bpe.py --encode unhappiness
Eleven pieces. The tokeniser has never seen anything like it and falls back to characters. That spread is fertility.
Download bpe.py
· served verbatim at https://nlp.jcrlabz.com/code/bpe.py
"""Byte pair encoding, traced step by step as Chapter 5 traces it.
Run this and you get the same corpus tables, the same pair counts, the same
ties and the same five merges the book prints.
python3 bpe.py # the low/lower/newest/widest corpus
python3 bpe.py --merges 8 # keep going and watch it continue
python3 bpe.py --encode lowest # apply the learned rules to a new word
python3 bpe.py --no-tiebreak # see why the tie-break rule exists
The two functions that matter are `train` and `encode`. Everything else is
printing.
Install: nothing, the Python standard library is enough
"""
import argparse
from collections import Counter
END = "</w>"
def pair_counts(words):
"""Every adjacent pair, weighted by the frequency of the word it sits in.
The weighting is the step most hand simulations get wrong. A pair inside a
word seen 6 times contributes 6, not 1.
"""
counts = Counter()
for symbols, freq in words.items():
for a, b in zip(symbols, symbols[1:]):
counts[(a, b)] += freq
return counts
def merge_pair(words, pair):
"""Rewrite the corpus with `pair` joined into one symbol."""
a, b = pair
out = {}
for symbols, freq in words.items():
joined, i = [], 0
while i < len(symbols):
if i + 1 < len(symbols) and symbols[i] == a and symbols[i + 1] == b:
joined.append(a + b)
i += 2
else:
joined.append(symbols[i])
i += 1
out[tuple(joined)] = out.get(tuple(joined), 0) + freq
return out
def train(corpus, num_merges, tiebreak=True, trace=None):
"""Learn an ordered merge list from {word: frequency}.
Returns the merges as "a b" strings, which is the form `encode` expects
and the form a GPT-2 merges.txt file uses.
"""
words = {tuple(list(w) + [END]): f for w, f in corpus.items()}
merges = []
for step in range(1, num_merges + 1):
counts = pair_counts(words)
if not counts:
break
top = max(counts.values())
tied = sorted(p for p, n in counts.items() if n == top)
# Ties are common in small corpora. Break them lexicographically or
# the run is not reproducible.
best = tied[0] if tiebreak else tied[-1]
if trace is not None:
trace(step, words, counts, top, tied, best)
merges.append(" ".join(best))
words = merge_pair(words, best)
return merges, words
def encode(word, merges):
"""Apply the merge rules IN LEARNED ORDER to one new word."""
symbols = list(word) + [END]
history = [("split", list(symbols), True)]
for rule in merges:
a, b = rule.split(" ")
before = list(symbols)
symbols = list(merge_pair({tuple(symbols): 1}, (a, b)).keys())[0]
history.append((rule, list(symbols), list(symbols) != before))
return list(symbols), history
# ----------------------------------------------------------------- printing
def show_corpus(words, title):
print(f"\n {title}")
print(f" {'symbols':<28}{'frequency'}")
for symbols, freq in words.items():
print(f" {' '.join(symbols):<28}{freq}")
def show_counts(counts, top, limit=7):
print(f"\n {'pair':<16}{'count'}")
for pair, n in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[:limit]:
flag = " <-- tied for first" if n == top else ""
print(f" {' '.join(pair):<16}{n}{flag}")
if len(counts) > limit:
print(f" ... and {len(counts) - limit} more, all lower")
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--merges", type=int, default=5)
ap.add_argument("--encode", default="lowest",
help="word to tokenise with the learned rules")
ap.add_argument("--no-tiebreak", action="store_true",
help="take the lexicographically LAST tied pair instead")
a = ap.parse_args()
corpus = {"low": 5, "lower": 2, "newest": 6, "widest": 3}
print("\nCORPUS")
print(f" {'word':<10}{'frequency'}")
for w, f in corpus.items():
print(f" {w:<10}{f}")
def trace(step, words, counts, top, tied, best):
print(f"\n{'=' * 62}\nSTEP {step}")
show_corpus(words, "corpus going in:")
show_counts(counts, top)
if len(tied) > 1:
names = ", ".join(f"'{x} {y}'" for x, y in tied)
print(f"\n {len(tied)} pairs tie at {top}: {names}")
print(f" tie-break takes '{best[0]} {best[1]}'")
else:
print(f"\n clear winner at {top}, no tie")
print(f"\n MERGE: '{best[0]} {best[1]}' -> {best[0] + best[1]}")
merges, final = train(corpus, a.merges,
tiebreak=not a.no_tiebreak, trace=trace)
print(f"\n{'=' * 62}")
show_corpus(final, "corpus after the last merge:")
print(f"\nTHE MERGE LIST (this ordered list IS the tokeniser)")
for i, m in enumerate(merges, 1):
a_, b_ = m.split(" ")
print(f" {i}. '{m}' -> {a_ + b_}")
print(f"\nTRAINING WORDS, TOKENISED")
for w in corpus:
toks, _ = encode(w, merges)
print(f" {w:<10}{' '.join(toks)}")
print(f"\nAN UNSEEN WORD: {a.encode!r}")
toks, history = encode(a.encode, merges)
for rule, state, changed in history:
label = "split" if rule == "split" else f"rule '{rule}'"
note = "" if changed or rule == "split" else " (no match)"
print(f" {label:<20}{' '.join(state)}{note}")
print(f"\n result: {' '.join(toks)}")
if __name__ == "__main__":
main()