Spaces:
Sleeping
Sleeping
| #!/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() | |