Initial commit

scripts/prepare_ablation_splits.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+"""Prepare fixed data splits for ablation experiments."""
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from pathlib import Path
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+DATASETS = {
+    "searchqa": {
+        "raw": PROJECT_ROOT / "data/searchqa_train_2000.json",
+        "out": PROJECT_ROOT / "data/ablation_splits/searchqa",
+        "filenames": {"train": "train.json", "val": "selection.json", "test": "test.json"},
+    },
+    "spreadsheetbench": {
+        "raw": PROJECT_ROOT / "data/spreadsheetbench_verified_400/dataset.json",
+        "out": PROJECT_ROOT / "data/ablation_splits/spreadsheetbench",
+        "filenames": {"train": "train.json", "val": "sel.json", "test": "test.json"},
+    },
+}
+
+SPLITS = ("1shot", "1:1:8", "2:1:7", "4:1:5")
+
+
+def load_items(path: Path) -> list[dict]:
+    with path.open(encoding="utf-8") as f:
+        data = json.load(f)
+    if not isinstance(data, list):
+        raise TypeError(f"Expected JSON array in {path}, got {type(data).__name__}")
+    return data
+
+
+def split_counts(total: int, split: str) -> tuple[int, int, int]:
+    if split == "1shot":
+        if total < 3:
+            raise ValueError(f"Need at least 3 items for 1shot split, got {total}")
+        return 1, 1, total - 2
+
+    ratio = split
+    weights = [int(part) for part in ratio.split(":")]
+    if len(weights) != 3 or min(weights) <= 0:
+        raise ValueError(f"Invalid ratio: {ratio}")
+    denom = sum(weights)
+    raw = [total * weight / denom for weight in weights]
+    counts = [int(value) for value in raw]
+    remaining = total - sum(counts)
+    order = sorted(
+        range(3),
+        key=lambda idx: (raw[idx] - counts[idx], weights[idx]),
+        reverse=True,
+    )
+    for idx in order[:remaining]:
+        counts[idx] += 1
+    return counts[0], counts[1], counts[2]
+
+
+def split_tag(split: str) -> str:
+    return "1shot" if split == "1shot" else split.replace(":", "-")
+
+
+def write_json(path: Path, items: list[dict]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(items, f, ensure_ascii=False, indent=2)
+
+
+def prepare_dataset(name: str, *, seed: int, force: bool) -> None:
+    spec = DATASETS[name]
+    raw_path = spec["raw"]
+    out_root = spec["out"]
+    filenames = spec["filenames"]
+
+    items = load_items(raw_path)
+    for split in SPLITS:
+        ratio_tag = split_tag(split)
+        split_dir = out_root / f"{ratio_tag}_seed{seed}"
+        manifest_path = split_dir / "split_manifest.json"
+        if manifest_path.exists() and not force:
+            print(f"skip {name} {split}: {split_dir} exists")
+            continue
+
+        shuffled = list(items)
+        random.Random(seed).shuffle(shuffled)
+        train_n, val_n, test_n = split_counts(len(shuffled), split)
+        train_items = shuffled[:train_n]
+        val_items = shuffled[train_n: train_n + val_n]
+        test_items = shuffled[train_n + val_n: train_n + val_n + test_n]
+
+        write_json(split_dir / "train" / filenames["train"], train_items)
+        write_json(split_dir / "val" / filenames["val"], val_items)
+        write_json(split_dir / "test" / filenames["test"], test_items)
+        write_json(
+            manifest_path,
+            {
+                "dataset": name,
+                "source": str(raw_path),
+                "split_mode": "precomputed_ratio",
+                "split_name": split,
+                "split_ratio": split if split != "1shot" else "1 train / 1 val / rest test",
+                "split_seed": seed,
+                "counts": {
+                    "train": len(train_items),
+                    "val": len(val_items),
+                    "test": len(test_items),
+                },
+            },
+        )
+        print(
+            f"wrote {name} {split} -> {split_dir} "
+            f"(train={len(train_items)}, val={len(val_items)}, test={len(test_items)})"
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--force", action="store_true")
+    parser.add_argument("--dataset", choices=sorted(DATASETS), action="append")
+    args = parser.parse_args()
+
+    for name in args.dataset or sorted(DATASETS):
+        prepare_dataset(name, seed=args.seed, force=args.force)
+
+
+if __name__ == "__main__":
+    main()