"""Quick screening set (274 dev items, never test): 100 math, 50 code, 70 MMLU-Pro (5/category), 54 BBH (2/task).
Random, NOT selected by the original model's results (that would bias every comparison by regression to the mean)."""
import json, random, collections
rng = random.Random(2024)
dev = [json.loads(l) for l in open("data/train/dev.jsonl", encoding="utf-8")]
gen = [json.loads(l) for l in open("data/train/dev_general.jsonl", encoding="utf-8")]
math = [r for r in dev if r["task"] == "math"]
code = [r for r in dev if r["task"] != "math"]
pick = rng.sample(math, 100) + rng.sample(code, 50)
for task, k in [("mmlupro", 5), ("bbh", 2)]:
    groups = collections.defaultdict(list)
    for r in gen:
        if r["task"] == task:
            groups[r.get("category") or r["id"].split("/")[1]].append(r)
    for g in sorted(groups):
        pick += rng.sample(groups[g], k)
with open("data/quick/quick.jsonl", "w", encoding="utf-8") as f:
    for r in pick:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")
print(collections.Counter(r["task"] for r in pick))