precompute_merges.py

"""Generate merges/data.json — every figure on /merges/, computed not written."""

import json
import os

import bpe
from cjsload import Reference
from corpora import BERRIES

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..",
                   "vocabulary", "data.json")

N_MERGES = 30

# Words that show the three outcomes: learned whole, learned in pieces, unseen.
PROBES = [" strawberry", "strawberry", " blackberry", "strawberries", " kiwi"]

# Where the assembly of " strawberry" visibly changes.
CHECKPOINTS = [0, 1, 2, 3, 4, 16, 17]


def main():
    trained = bpe.train(BERRIES, N_MERGES)
    merges = trained["merges"]
    real = Reference("o200k_base")

    data = {
        "corpus": BERRIES,
        "corpus_chars": len(BERRIES),
        "corpus_words": len(bpe.pretokenize(BERRIES)),
        "corpus_unique": len(set(bpe.pretokenize(BERRIES))),
        "alphabet": trained["alphabet"],
        "steps": trained["steps"],
        "vocab_size": trained["vocab_size"],
        "n_merges": len(merges),
        "checkpoints": [
            {"after": n, "symbols": bpe.encode(" strawberry", merges[:n])}
            for n in CHECKPOINTS
        ],
        "probes": [
            {
                "text": p,
                "toy": bpe.encode(p, merges),
                "real": [t["text"] for t in real.tokens(p)],
            }
            for p in PROBES
        ],
    }

    with open(OUT, "w", encoding="utf-8") as fh:
        json.dump(data, fh, ensure_ascii=False, indent=1)

    print(f"wrote {os.path.relpath(OUT, os.path.dirname(OUT))}")
    print(f"  corpus: {data['corpus_chars']} chars, {data['corpus_words']} words, "
          f"{data['corpus_unique']} unique")
    print(f"  {data['n_merges']} merges, alphabet {len(data['alphabet'])}, "
          f"vocab {data['vocab_size']}")
    print("  first merges: " +
          " ".join(f"{s['token']!r}" for s in data["steps"][:5]))
    for p in data["probes"]:
        print(f"  {p['text']!r:>16}  toy {len(p['toy']):>2}  real {len(p['real']):>2}")


if __name__ == "__main__":
    main()