inference.py

"""
Inference only. No training, no evaluation harness.

This is the file to read if you just want to run the model, and the file the
Hugging Face repo is built around.

    python inference.py                                  # write 5 queries
    python inference.py --prompt "SELECT region ,"       # complete a prompt
    python inference.py --n 10 --temperature 1.2         # sample more, wilder
    python inference.py --run                            # execute them for real
    python inference.py --probs "SELECT region , SUM ( qty ) FROM sales GROUP BY"

    python inference.py --export hf/tiny-sql-gpt         # package for the Hub

Programmatic use:

    from inference import TinySQLGPT
    m = TinySQLGPT.from_pretrained("checkpoints/tiny.pt")
    print(m.generate())
"""

import argparse
import json
import os

import torch
from torch.nn import functional as F

from tiny_gpt import BOS, Config, Tokenizer, TinyGPT, HERE

DEFAULT_CKPT = os.path.join(HERE, "checkpoints", "tiny.pt")


class TinySQLGPT:
    """A trained model plus its tokenizer. Nothing else."""

    def __init__(self, model, tok, device="cpu"):
        self.model, self.tok, self.device = model, tok, device

    # ── loading ──────────────────────────────────────────────────────────────

    @classmethod
    def from_pretrained(cls, path=DEFAULT_CKPT, device="cpu"):
        """Accepts a .pt checkpoint, a directory in Hub layout, or a repo id."""
        if not os.path.exists(path) and "/" in path and not path.endswith(".pt"):
            from huggingface_hub import snapshot_download   # optional dependency
            path = snapshot_download(repo_id=path)

        if os.path.isdir(path):
            return cls._from_dir(path, device)
        return cls._from_ckpt(path, device)

    @classmethod
    def _from_ckpt(cls, path, device):
        if not os.path.exists(path):
            raise SystemExit(f"No checkpoint at {path}. Run: python tiny_gpt.py --train")
        ck = torch.load(path, map_location=device, weights_only=False)
        return cls._build(ck["cfg"], ck["state_dict"], ck["itos"], device)

    @classmethod
    def _from_dir(cls, path, device):
        """Hub layout: config.json + weights, no pickled Python objects."""
        with open(os.path.join(path, "config.json")) as f:
            meta = json.load(f)
        safe = os.path.join(path, "model.safetensors")
        if os.path.exists(safe):
            from safetensors.torch import load_file
            state = load_file(safe)
        else:
            state = torch.load(os.path.join(path, "pytorch_model.bin"),
                               map_location=device, weights_only=True)
        return cls._build(meta["cfg"], state, meta["itos"], device)

    @classmethod
    def _build(cls, cfg_dict, state, itos, device):
        model = TinyGPT(Config(**cfg_dict)).to(device)
        model.load_state_dict(state)
        model.eval()
        tok = Tokenizer.__new__(Tokenizer)
        tok.itos = itos
        tok.stoi = {s: i for i, s in enumerate(itos)}
        return cls(model, tok, device)

    # ── generation ───────────────────────────────────────────────────────────

    def _ids(self, prompt):
        ids = [self.tok.stoi[BOS]]
        if prompt:
            ids += self.tok.encode(prompt)
        return torch.tensor([ids], dtype=torch.long, device=self.device)

    def generate(self, prompt="", temperature=0.8, top_k=None, max_new_tokens=None):
        """Return one SQL query as a string."""
        idx = self._ids(prompt)
        room = self.model.cfg.block_size - idx.shape[1]
        out = self.model.generate(
            idx, max_new_tokens=max_new_tokens or room,
            temperature=temperature, top_k=top_k, stop=self.tok.stoi[";"])
        return self.tok.decode(out[0, 1:].tolist())

    def generate_many(self, n=5, **kw):
        return [self.generate(**kw) for _ in range(n)]

    def next_token_probs(self, prompt, top=10):
        """What the model thinks comes next, as (token, probability) pairs.

        The whole distribution is only 155 wide, so `top=None` really does
        return all of it, the thing you cannot do with a frontier model.
        """
        with torch.no_grad():
            logits, _ = self.model(self._ids(prompt))
        probs = F.softmax(logits[0, -1], dim=-1)
        k = top or len(self.tok)
        vals, idx = torch.topk(probs, min(k, len(self.tok)))
        return [(self.tok.itos[i], float(p)) for p, i in zip(vals, idx)]

    @property
    def n_params(self):
        return self.model.n_params()

    # ── export ───────────────────────────────────────────────────────────────

    def export(self, outdir):
        """Write a Hugging Face style folder: config.json + weights.

        Deliberately avoids a pickled checkpoint. Nobody should have to run
        torch.load(weights_only=False) on a stranger's file.
        """
        os.makedirs(outdir, exist_ok=True)
        with open(os.path.join(outdir, "config.json"), "w") as f:
            json.dump({
                "model_type": "tiny-sql-gpt",
                "cfg": {k: v for k, v in vars(self.model.cfg).items()},
                "itos": self.tok.itos,
            }, f, indent=2)
        state = {k: v.contiguous() for k, v in self.model.state_dict().items()}
        try:
            from safetensors.torch import save_file
            save_file(state, os.path.join(outdir, "model.safetensors"))
            wrote = "model.safetensors"
        except ImportError:
            torch.save(state, os.path.join(outdir, "pytorch_model.bin"))
            wrote = "pytorch_model.bin"
        return wrote


def main():
    ap = argparse.ArgumentParser(description="Tiny SQL GPT inference")
    ap.add_argument("--model", default=DEFAULT_CKPT,
                    help=".pt file, Hub-layout directory, or HF repo id")
    ap.add_argument("--prompt", default="", help="text to continue")
    ap.add_argument("--n", type=int, default=5)
    ap.add_argument("--temperature", type=float, default=0.8)
    ap.add_argument("--top-k", type=int, default=None)
    ap.add_argument("--run", action="store_true",
                    help="execute each query against SQLite and report")
    ap.add_argument("--probs", metavar="PROMPT",
                    help="show the next-token distribution for a prompt")
    ap.add_argument("--export", metavar="DIR", help="write a Hub-ready folder")
    ap.add_argument("--device", default="cpu")
    args = ap.parse_args()

    m = TinySQLGPT.from_pretrained(args.model, args.device)

    if args.export:
        wrote = m.export(args.export)
        print(f"exported to {args.export}/ (config.json + {wrote})")
        return

    if args.probs:
        print(f"context: {args.probs}\n")
        for t, p in m.next_token_probs(args.probs):
            print(f"  {t:<12} {p:6.3f}  {'█' * int(round(p * 30))}")
        return

    queries = m.generate_many(args.n, prompt=args.prompt,
                              temperature=args.temperature, top_k=args.top_k)
    if not args.run:
        print("\n".join(queries))
        return

    import evaluate                      # only needed for --run
    conn = evaluate.build_db()
    ok = 0
    for q in queries:
        good = evaluate.executes(conn, q)
        ok += good
        print(f"  [{'OK  ' if good else 'FAIL'}] {q}")
    print(f"\n{ok}/{len(queries)} executed against a real database "
          f"({m.n_params:,} parameters)")


if __name__ == "__main__":
    main()