File size: 1,221 Bytes
e379257 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | """Write a JSON record of a train/eval split: ids and per-scenario counts per file, plus which scenarios are missing.
usage: python scripts/record_split.py out.json name1=file1.jsonl name2=file2.jsonl ...
"""
import json, sys
from collections import Counter
out, specs = sys.argv[1], sys.argv[2:]
res = {}
for s in specs:
name, path = s.split("=", 1)
with open(path, encoding="utf-8") as f:
rows = [json.loads(l) for l in f if l.strip()]
sc = Counter((r.get("meta") or {}).get("scenario") or "?" for r in rows)
res[name] = {"path": path, "n": len(rows), "n_with_box_questions": sum(any(q.get("type") == "box" for q in r["questions"].values()) for r in rows),
"scenarios": dict(sorted(sc.items())), "ids": [r["id"] for r in rows]}
allsc = set().union(*(set(v["scenarios"]) for v in res.values()))
for v in res.values():
v["missing_scenarios"] = sorted(allsc - set(v["scenarios"]))
with open(out, "w", encoding="utf-8", newline="\n") as f:
json.dump(res, f, indent=1, ensure_ascii=False)
print(json.dumps({k: {"n": v["n"], "box": v["n_with_box_questions"], "n_scen": len(v["scenarios"]), "missing": v["missing_scenarios"]} for k, v in res.items()}))
|