capture.py

"""Capture real outputs so the static page can show them without a server.

A static Space has no backend: nothing can run GPT-2 for the visitor. What it
can do is show what GPT-2 really produced, side by side, and say plainly that
these are recorded runs anyone can reproduce with `pip install sillage`.

Everything written to data.json here comes out of the tool, unedited.

    python capture.py     ->  data.json
"""

import json
import os
import tempfile

import sample
from sillage import Sillage

HERE = os.path.dirname(os.path.abspath(__file__))
STATE = os.path.join(HERE, "state")
N = 20

PROMPTS = [
    "On a 36k-token stream of novel technical text, the memory",
    "the memory improves GPT-2's test negative log-likelihood by",
    "At 500k tokens the fixed matrix",
    "surprise gating quadruples the gain of",
    "Every write is gated by the model's own",
    "a three-factor plasticity rule whose modulator is",
]


def main():
    memory = Sillage(model="gpt2", state=STATE, quiet=False)
    frozen = Sillage(model="gpt2", state=tempfile.mkdtemp(), quiet=True)
    tok, model = memory.load_model()
    frozen._tok, frozen._model = tok, model

    out = {"model": "openai-community/gpt2",
           "read": "papers/sillage/sillage.tex (8969 tokens)",
           "state_mb": round(sum(
               os.path.getsize(os.path.join(STATE, f))
               for f in os.listdir(STATE)) / 1e6, 1),
           "completions": []}
    for p in PROMPTS:
        a, b = frozen.complete(p, n=N), memory.complete(p, n=N)
        out["completions"].append({"prompt": p, "frozen": a, "memory": b,
                                   "same": a.strip() == b.strip()})
        print(f"  {p[:44]:46s} {'=' if a == b else 'differs'}")

    import app                      # reuse the demo's instrumented read loop
    plot, summary, fixes, _, suggestion = app.read_and_report(sample.MANUAL)
    out["manual"] = {"summary": summary, "fixes": fixes or [],
                     "perplexity": plot, "suggestion": suggestion}
    a, b = frozen.complete(suggestion, n=12), None
    fresh = Sillage(model="gpt2", state=app.SESSIONS[-1], quiet=True)
    fresh._tok, fresh._model = tok, model
    b = fresh.complete(suggestion, n=12)
    out["manual"]["completion"] = {"prompt": suggestion, "frozen": a,
                                   "memory": b}

    with open(os.path.join(HERE, "data.json"), "w", encoding="utf-8") as f:
        json.dump(out, f, indent=1, ensure_ascii=False)
    print(f"wrote data.json: {len(out['completions'])} completions, "
          f"{len(out['manual']['fixes'])} corrections")


if __name__ == "__main__":
    main()