| """MBPP sanitized load/filter for small-model coding RL."""
|
| from __future__ import annotations
|
|
|
| import json
|
| from pathlib import Path
|
| from typing import Any
|
|
|
| from deepagents_rl.rewards.unit_test import run_hidden_tests
|
|
|
|
|
| def _row_to_task(ex: dict[str, Any], split: str) -> dict[str, Any] | None:
|
| """Normalize an MBPP sanitized example."""
|
| task_id = ex.get("task_id") or ex.get("id")
|
| text = ex.get("text") or ex.get("prompt") or ex.get("description") or ""
|
| code = ex.get("code") or ex.get("reference_code") or ""
|
| test_list = list(ex.get("test_list") or ex.get("test_cases") or [])
|
| if not text or not test_list:
|
| return None
|
|
|
| visible = [test_list[0]] if test_list else []
|
| hidden = list(test_list)
|
| prompt = (
|
| f"Write a Python solution for the following problem.\n\n"
|
| f"{text.strip()}\n\n"
|
| f"Visible example test (not the full grader):\n{visible[0]}\n\n"
|
| f"Return ONLY a ```python``` code block with the full solution."
|
| )
|
| return {
|
| "task_id": str(task_id),
|
| "split": split,
|
| "text": text.strip(),
|
| "reference_code": code,
|
| "visible_tests": visible,
|
| "hidden_tests": hidden,
|
| "prompt": prompt,
|
| "test_list": test_list,
|
| }
|
|
|
|
|
| def filter_solvable(tasks: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
| """Keep tasks whose reference solution passes hidden tests in our sandbox."""
|
| keep: list[dict[str, Any]] = []
|
| drop: list[dict[str, Any]] = []
|
| for t in tasks:
|
| ref = t.get("reference_code") or ""
|
| if not ref.strip():
|
| drop.append({**t, "drop_reason": "no_reference"})
|
| continue
|
| grade = run_hidden_tests(ref, t["hidden_tests"], timeout_s=5.0)
|
| if grade["pass_frac"] >= 1.0:
|
| keep.append(t)
|
| else:
|
| drop.append({**t, "drop_reason": "ref_fail", "grade": grade})
|
| return keep, drop
|
|
|
|
|
| def load_mbpp_sanitized() -> dict[str, list[dict[str, Any]]]:
|
| """Load google-research-datasets/mbpp sanitized config from HF."""
|
| from datasets import load_dataset
|
|
|
|
|
| try:
|
| ds = load_dataset("google-research-datasets/mbpp", "sanitized")
|
| except Exception:
|
| ds = load_dataset("mbpp", "sanitized")
|
|
|
| out: dict[str, list[dict[str, Any]]] = {}
|
| for split in ds:
|
| tasks = []
|
| for ex in ds[split]:
|
| t = _row_to_task(dict(ex), split=split)
|
| if t:
|
| tasks.append(t)
|
| out[split] = tasks
|
| return out
|
|
|
|
|
| def build_train_eval(
|
| *,
|
| filter_refs: bool = True,
|
| ) -> dict[str, list[dict[str, Any]]]:
|
| """
|
| Train = sanitized train + validation + prompt (none of test).
|
| Eval = sanitized test (held-out).
|
| """
|
| raw = load_mbpp_sanitized()
|
| train_parts = []
|
| for key in ("train", "validation", "prompt"):
|
| train_parts.extend(raw.get(key, []))
|
| eval_parts = list(raw.get("test", []))
|
|
|
| drop_log: list[dict[str, Any]] = []
|
| if filter_refs:
|
| train_parts, d1 = filter_solvable(train_parts)
|
| eval_parts, d2 = filter_solvable(eval_parts)
|
| drop_log = d1 + d2
|
|
|
|
|
| train_ids = {t["task_id"] for t in train_parts}
|
| eval_parts = [t for t in eval_parts if t["task_id"] not in train_ids]
|
|
|
| return {
|
| "train": train_parts,
|
| "eval": eval_parts,
|
| "dropped": drop_log,
|
| }
|
|
|
|
|
| def save_jsonl(tasks: list[dict[str, Any]], path: Path) -> None:
|
| path.parent.mkdir(parents=True, exist_ok=True)
|
| with path.open("w", encoding="utf-8") as f:
|
| for t in tasks:
|
|
|
| row = {k: v for k, v in t.items() if k not in ("drop_reason", "grade")}
|
| f.write(json.dumps(row) + "\n")
|
|
|
|
|
| def to_hf_dataset(splits: dict[str, list[dict[str, Any]]]):
|
| from datasets import Dataset, DatasetDict
|
|
|
| def _clean(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
| cleaned = []
|
| for t in rows:
|
| cleaned.append(
|
| {
|
| "task_id": t["task_id"],
|
| "split": t["split"],
|
| "text": t["text"],
|
| "prompt": t["prompt"],
|
| "reference_code": t.get("reference_code", ""),
|
| "visible_tests": t["visible_tests"],
|
| "hidden_tests": t["hidden_tests"],
|
| }
|
| )
|
| return cleaned
|
|
|
| return DatasetDict(
|
| {
|
| "train": Dataset.from_list(_clean(splits["train"])),
|
| "eval": Dataset.from_list(_clean(splits["eval"])),
|
| }
|
| )
|
|
|
