tools/build_space_data.py

#!/usr/bin/env python3
"""Package the precompute cache into the small files the Space reads.

    python hf_space_demo/tools/build_space_data.py \
        --cache-dir /scratch/lhotsko/C2L_SPACE_CACHE --out-dir hf_space_demo/data

Writes ``index.json`` plus, per selected repository, ``<safe>.json.gz``
(commits + assertions + cached predictions) and ``<safe>.ctx.npy``
(float32 [n_commits, 2048] GRU contexts).
"""

from __future__ import annotations

import argparse
import gzip
import json
import pickle
import urllib.request
from pathlib import Path

import numpy as np


def github_meta(repo_id: str) -> dict:
    try:
        with urllib.request.urlopen(f"https://api.github.com/repos/{repo_id}", timeout=10) as r:
            d = json.load(r)
        return {"description": d.get("description") or "", "stars": d.get("stargazers_count")}
    except Exception:
        return {"description": "", "stars": None}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache-dir", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--min-commits", type=int, default=10)
    ap.add_argument("--max-commits", type=int, default=400)
    ap.add_argument("--max-repos", type=int, default=24)
    args = ap.parse_args()

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    candidates = []
    for p in sorted(Path(args.cache_dir).glob("*.pkl")):
        with open(p, "rb") as f:
            d = pickle.load(f)
        n_c, qn = len(d["commits"]), d["qnas"]
        if not (args.min_commits <= n_c <= args.max_commits) or not qn:
            continue
        base = sum(q["base_em"] for q in qn) / len(qn)
        c2l = sum(q["c2l_em"] for q in qn) / len(qn)
        tt = sum(1 for q in qn if q.get("prev_em") == 0.0)
        candidates.append((c2l - base, tt, p, d, base, c2l))

    candidates.sort(key=lambda x: (x[0], x[1]), reverse=True)
    index = []
    for gain, tt, p, d, base, c2l in candidates[: args.max_repos]:
        safe = p.stem
        commits = d["commits"]
        per_commit = [[] for _ in commits]
        for q in d["qnas"]:
            per_commit[q["k"]].append(q)
        for k, c in enumerate(commits):
            qs = per_commit[k]
            c["n_qnas"] = len(qs)
            c["base_em"] = sum(q["base_em"] for q in qs) / len(qs) if qs else None
            c["c2l_em"] = sum(q["c2l_em"] for q in qs) / len(qs) if qs else None
            c["time_travel"] = any(q.get("prev_em") == 0.0 for q in qs)
        payload = {"repo_id": d["repo_id"], "commits": commits, "qnas": per_commit,
                   "max_input_tokens": d["max_input_tokens"],
                   "max_new_tokens": d["max_new_tokens"]}
        with gzip.open(out / f"{safe}.json.gz", "wt", encoding="utf-8") as f:
            json.dump(payload, f)
        np.save(out / f"{safe}.ctx.npy", d["ctx"].astype(np.float32))
        dates = [c.get("timestamp") for c in commits if c.get("timestamp")]
        index.append({
            "repo_id": d["repo_id"], "safe": safe, "n_commits": len(commits),
            "n_qnas": len(d["qnas"]), "base_em": base, "c2l_em": c2l,
            "n_time_travel": tt, "first_date": min(dates) if dates else "",
            "last_date": max(dates) if dates else "", **github_meta(d["repo_id"]),
        })
        print(f"{d['repo_id']:45s} commits={len(commits):4d} base={base:.3f} "
              f"c2l={c2l:.3f} gain={gain:+.3f} time_travel={tt}")

    index.sort(key=lambda r: r["repo_id"].lower())
    (out / "index.json").write_text(json.dumps(index, indent=1))
    print(f"wrote {len(index)} repos to {out}")


if __name__ == "__main__":
    main()