File size: 4,875 Bytes
a31a731 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """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
# Leakage-aware: first assert may appear in prompt as visible example
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
# Prefer sanitized; fall back to full mbpp
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
# Disjoint task ids
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:
# Don't need full drop grades in published rows
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"])),
}
)
|