precompute_predict.py

"""Generate predict/data.json — every figure on /predict/."""

import json
import os

import ngram
from corpora import TALE

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

ORDER = 2
CONTEXT = "it was the"
SEED = 11
LENGTH = 22


def marked(full, kept):
    """Full distribution with a flag for whether each token survived the cut."""
    keep = {t for t, _ in kept}
    reshaped = dict(kept)
    return [{"token": t, "p": p, "kept": t in keep, "after": reshaped.get(t, 0.0)}
            for t, p in full]


def main():
    trained = ngram.train(TALE, ORDER)
    context = ngram.tokenize(CONTEXT)
    items, used = ngram.distribution(trained, context)

    data = {
        "corpus": TALE,
        "corpus_tokens": len(trained["tokens"]),
        "corpus_vocab": len(trained["vocab"]),
        "order": ORDER,
        "context": CONTEXT,
        "context_order_used": used,
        "distribution": [{"token": t, "p": p} for t, p in items],
        "temperatures": [
            {"temp": temp,
             "items": [{"token": t, "p": p}
                       for t, p in ngram.apply_temperature(items, temp)]}
            for temp in [0.0, 0.5, 1.0, 2.0]
        ],
        "top_k": {"k": 3, "items": marked(items, ngram.apply_top_k(items, 3))},
        "top_p": {"p": 0.5, "items": marked(items, ngram.apply_top_p(items, 0.5))},
        "samples": [
            {"label": label, "temp": temp, "top_k": k, "top_p": p,
             "text": ngram.generate(trained, CONTEXT, LENGTH, temp, k, p,
                                    SEED)["text"]}
            for label, temp, k, p in [
                ("Greedy (T = 0)", 0.0, 0, 0.0),
                ("T = 0.7", 0.7, 0, 0.0),
                ("T = 1.0", 1.0, 0, 0.0),
                ("T = 2.5", 2.5, 0, 0.0),
            ]
        ],
        "orders": [
            {"order": o,
             "text": ngram.generate(ngram.train(TALE, o), "it was", 18, 1.0,
                                    0, 0.0, 5)["text"]}
            for o in [2, 1, 0]
        ],
        "seed": SEED,
    }

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

    print(f"wrote predict/data.json")
    print(f"  corpus {data['corpus_tokens']} tokens, vocab {data['corpus_vocab']}")
    print(f"  after {CONTEXT!r} (order {used}): "
          f"{len(data['distribution'])} candidates")
    for s in data["samples"]:
        print(f"  {s['label']:<16} {s['text'][:56]!r}")


if __name__ == "__main__":
    main()