code/scripts/demo.py

"""Watch the model work: it gets a freshly generated fake project (configs, code, tests, docs, logs,
CSVs) in a sandbox directory and a question about it, and solves it with bash/read/edit/write/submit.

  python scripts/demo.py --model rl                       # random task kind
  python scripts/demo.py --model rl --kind fix_test --seed 7
  python scripts/demo.py --model rl --question "Which service has the most replicas?"

--model is a directory with model.safetensors + config.json (e.g. the `rl` or `base` folder of
the Hugging Face repo) or a training checkpoint (.pt). Runs on CPU, Intel XPU or CUDA.
"""
import argparse
import dataclasses
import json
import os
import random
import sys

import torch
from tokenizers import Tokenizer

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tiny_agent.checkpoint import load_model  # noqa: E402
from tiny_agent.rollout import make_roller  # noqa: E402
from tiny_agent.tasks import KINDS, make_task  # noqa: E402


def device_auto():
    if hasattr(torch, "xpu") and torch.xpu.is_available():
        return "xpu"
    return "cuda" if torch.cuda.is_available() else "cpu"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="rl")
    ap.add_argument("--tokenizer", default="tokenizer.json")
    ap.add_argument("--kind", default=None, help=f"one of {', '.join(KINDS)}")
    ap.add_argument("--question", default=None, help="ask your own question about the generated project")
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--temperature", type=float, default=0.7)
    ap.add_argument("--device", default=None)
    a = ap.parse_args()
    dev = a.device or device_auto()
    seed = a.seed if a.seed is not None else random.randrange(10**6)
    rng = random.Random(seed)
    task = make_task(rng, a.kind or rng.choice(KINDS))
    if a.question:
        task = dataclasses.replace(task, question=a.question, answer="?")
    print(f"seed {seed} | kind {task.kind} | device {dev}\nproject files: {', '.join(sorted(task.files))}\n")
    print(f"USER: {task.question}\n")
    tok = Tokenizer.from_file(a.tokenizer)
    model = load_model(a.model, device=dev, dtype=torch.bfloat16 if dev != "cpu" else torch.float32).eval()
    roller = make_roller("lockstep", model, tok, device=dev, max_len=4096, temperature=a.temperature)
    ep = roller.run([task])[0]
    for m in ep.messages[2:]:
        if m["role"] == "assistant":
            if m.get("think"):
                print(f"  think: {m['think']}")
            if m.get("content"):
                print(f"  say:   {m['content']}")
            for c in m.get("tool_calls") or []:
                print(f"  call:  {json.dumps(c)}")
        elif m["role"] == "tool":
            for r in m["results"]:
                body = r if len(r) < 600 else r[:600] + f"\n  ... ({len(r)} chars)"
                print("  result: " + body.replace("\n", "\n          "))
        print()
    if a.question:
        print(f"submitted: {ep.ws.submitted!r}")
    else:
        print(f"submitted: {ep.ws.submitted!r} | expected: {task.answer!r} | correct: {ep.correct}")


if __name__ == "__main__":
    main()