scripts/capture_gallery.py

"""Capture one trace per tab into data/gallery/<paradigm>.json (shown when the page opens).

Usage:
  python scripts/capture_gallery.py --space USER/beyond-tokens-decoding --family qwen3
  python scripts/capture_gallery.py --url http://127.0.0.1:7860 --family qwen25
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from content import examples as EX  # noqa: E402

GALLERY = {
    "contrastive": {**EX.CONTRASTIVE[0], "max_new_tokens": 48},
    "guided": {**EX.GUIDED[0], "max_new_tokens": 40},
    "parallel": {**EX.PARALLEL[0], "max_new_tokens": 96},
}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url")
    ap.add_argument("--space")
    ap.add_argument("--family", default="qwen3")
    args = ap.parse_args()

    from gradio_client import Client

    from huggingface_hub import get_token

    token = os.environ.get("HF_TOKEN") or get_token()  # logged-in calls get the account's GPU quota
    client = Client(args.space, token=token) if args.space else Client(args.url or "http://127.0.0.1:7860")
    out = ROOT / "data" / "gallery"
    out.mkdir(parents=True, exist_ok=True)
    for paradigm, params in GALLERY.items():
        P = {k: v for k, v in params.items() if k != "label"}
        trace = client.predict(paradigm, args.family, json.dumps(P), api_name="/trace")
        path = out / f"{paradigm}.json"
        path.write_text(json.dumps(trace, ensure_ascii=False), encoding="utf-8")
        print(f"{paradigm:11} {params['label']:45} -> {path.relative_to(ROOT)}  {trace['runs']['method']['text'][:60]!r}")
    return 0


if __name__ == "__main__":
    sys.exit(main())