fix: filter mbpp+ overlap instead of aborting; mbpp+ spans multiple hf splits

Author: Devesh-Maheshwari

scripts/build_sft_dataset.py

@@ -93,24 +93,39 @@ def _build_prompt(text: str, test_list: list[str]) -> str:
 
 
 def load_mbpp_train() -> list[Task]:
-    """Return the MBPP train split as `Task` objects (skips ones we can't parse)."""
-    ds = cast("Any", load_dataset("mbpp", split="train"))
+    """Return MBPP tasks disjoint from MBPP+ (our held-out eval set).
+
+    EvalPlus's MBPP+ draws task_ids from across the whole of MBPP, not just
+    the "test" split — HuggingFace's `mbpp[train]` alone overlaps MBPP+ by
+    ~107 tasks. Strategy: load train + validation + prompt (skip HF's
+    "test" since MBPP+ is sourced from it), then filter out anything whose
+    task_id appears in MBPP+.
+    """
+    plus_ids: set[int] = {
+        int(t.task_id.split("/")[-1]) for t in load_mbpp_plus()
+    }
+
     tasks: list[Task] = []
-    for item in ds:
-        item_d = cast("dict[str, Any]", item)
-        test_list: list[str] = list(item_d.get("test_list") or [])
-        entry_point = _infer_entry_point(test_list)
-        if not entry_point or not test_list:
-            continue
-        tasks.append(
-            Task(
-                task_id=f"Mbpp/{item_d['task_id']}",
-                prompt=_build_prompt(str(item_d["text"]), test_list),
-                canonical_solution=str(item_d["code"]),
-                test="\n".join(test_list),
-                entry_point=entry_point,
+    for split in ("train", "validation", "prompt"):
+        ds = cast("Any", load_dataset("mbpp", split=split))
+        for item in ds:
+            item_d = cast("dict[str, Any]", item)
+            task_id_int = int(item_d["task_id"])
+            if task_id_int in plus_ids:
+                continue  # would leak into eval — skip
+            test_list: list[str] = list(item_d.get("test_list") or [])
+            entry_point = _infer_entry_point(test_list)
+            if not entry_point or not test_list:
+                continue
+            tasks.append(
+                Task(
+                    task_id=f"Mbpp/{task_id_int}",
+                    prompt=_build_prompt(str(item_d["text"]), test_list),
+                    canonical_solution=str(item_d["code"]),
+                    test="\n".join(test_list),
+                    entry_point=entry_point,
+                )
             )
-        )
     return tasks