export_web.py

"""
Export model.npz for the browser demo

Writes docs/model.js, which the demo page (docs/demo.html) loads to run
TinyGPT in the browser. Only the parameters are exported (no Adam state),
as float32, so the file stays small.

Usage:

    python export_web.py
    python export_web.py --model model.npz --out docs/model.js

Run it again after retraining so the demo uses the new weights.
"""

import argparse
import base64
import json
from pathlib import Path

import numpy as np

from tinygpt import load_model


# ====================================================================
# Export
#
# Each parameter becomes
#
#   { "shape": [rows, cols], "b64": base64(float32, row-major) }
#
# and the file assigns everything to window.TINYGPT_MODEL,
# so it loads with a plain <script> tag (works from file:// too).
# ====================================================================

def export(
    model_path,
    out_path
):

    (
        model,
        token_to_id,
        id_to_token,
        sentences,
        epoch
    ) = load_model(
        model_path
    )

    params = {}

    for k, v in model.p.items():

        a = np.ascontiguousarray(
            v,
            dtype="<f4"
        )

        params[k] = {
            "shape": list(a.shape),
            "b64": base64.b64encode(a.tobytes()).decode("ascii")
        }

    data = {
        "epoch": epoch,
        "config": {
            "vocab_size": model.V,
            "max_context": model.max_context,
            "d_model": model.d,
            "n_heads": model.H,
            "d_ff": model.df,
            "n_layers": model.N,
            "eps": model.eps
        },
        "vocab": [
            id_to_token[i]
            for i in range(model.V)
        ],
        "sentences": sentences,
        "params": params
    }

    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    out_path.write_text(
        "// Generated by export_web.py from "
        f"{Path(model_path).name}. Do not edit.\n"
        "window.TINYGPT_MODEL = "
        + json.dumps(data, separators=(",", ":"))
        + ";\n",
        encoding="utf-8"
    )

    n_params = sum(v.size for v in model.p.values())

    print(
        f"{out_path}: {n_params} parameters, "
        f"{out_path.stat().st_size / 1024:.0f} KB"
    )


if __name__ == "__main__":

    here = Path(__file__).parent

    parser = argparse.ArgumentParser(
        description="Export model.npz for the browser demo"
    )

    parser.add_argument(
        "--model",
        type=Path,
        default=here / "model.npz"
    )

    parser.add_argument(
        "--out",
        type=Path,
        default=here / "docs" / "model.js"
    )

    args = parser.parse_args()

    export(
        args.model,
        args.out
    )