Generate jsonl ists evaluation data

Author: jvamvas

experiments/scripts/save_eval_dataset.py

@@ -0,0 +1,42 @@
+import argparse
+from pathlib import Path
+
+import jsonlines
+
+from rsd.experiments.utils import load_summary_benchmarks
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--split", type=str, default="test", help="Dataset split (train or test)", choices=["train", "test"])
+parser.add_argument("--limit-samples", type=int, default=None)
+args = parser.parse_args()
+
+out_dir = Path(__file__).parent.parent.parent.parent / "ists_evaluation_data" / args.split
+out_dir.mkdir(parents=True, exist_ok=True)
+out_file = out_dir / "rsd.jsonl"
+
+benchmark_names = [
+    "ists",
+    "ists_negatives",
+    "ists_documents",
+    "ists_permutations",
+    "ists_de",
+    "ists_es",
+    "ists_fr",
+    "ists_ja",
+    "ists_ko",
+    "ists_zh",
+]
+
+with jsonlines.open(out_file, "w") as f:
+    benchmarks = load_summary_benchmarks(args.split)
+    results = {}
+    for j, benchmark in enumerate(benchmarks):
+        dataset = benchmark.to_dataset()
+        dataset.shuffle(seed=42)
+        if args.limit_samples:
+            dataset = dataset.select(range(args.limit_samples))
+        for i, sample in enumerate(dataset):
+            sample = dict(sample)
+            sample["subset"] = benchmark_names[j]
+            sample["id"] = f"{benchmark_names[j]}_{i}"
+            f.write(sample)