| |
| """Convert our own solved R2E rollouts to pi-format SFT. Not eval-suite items.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| from datasets import Dataset |
|
|
| OUT = Path("/var/lib/agentptb-cache/tb1/tmp/sft-data/pi-sft-solved2") |
| ROOTS = [ |
| Path("/mnt/pvc/users/simon/agentptb/runs/grok/workspace/outputs/rl-r2e4/run_default/rollouts"), |
| Path("/mnt/pvc/users/simon/agentptb/runs/grok/workspace/outputs/rl-r2e7/run_default/rollouts"), |
| Path("/mnt/pvc/users/simon/agentptb/runs/grok/workspace/outputs/rl-r2e8/run_default/rollouts"), |
| ] |
| TOOLS = json.dumps( |
| [ |
| { |
| "type": "function", |
| "function": { |
| "name": "bash", |
| "description": "Execute a bash command", |
| "parameters": { |
| "type": "object", |
| "properties": {"command": {"type": "string"}}, |
| "required": ["command"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "read", |
| "description": "Read a file", |
| "parameters": { |
| "type": "object", |
| "properties": {"path": {"type": "string"}}, |
| "required": ["path"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "write", |
| "description": "Write a file", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "path": {"type": "string"}, |
| "content": {"type": "string"}, |
| }, |
| "required": ["path", "content"], |
| }, |
| }, |
| }, |
| { |
| "type": "function", |
| "function": { |
| "name": "edit", |
| "description": "Edit a file", |
| "parameters": { |
| "type": "object", |
| "properties": {"path": {"type": "string"}, "edits": {"type": "array"}}, |
| "required": ["path", "edits"], |
| }, |
| }, |
| }, |
| ] |
| ) |
|
|
|
|
| def score(tr, key): |
| v = (tr.get("rewards") or {}).get(key) |
| return float((v.get("score") if isinstance(v, dict) else v) or 0) |
|
|
|
|
| def norm_tc(tc: dict) -> dict: |
| if "function" in tc: |
| return { |
| "id": tc.get("id") or "call_0", |
| "type": "function", |
| "function": { |
| "name": (tc.get("function") or {}).get("name") or "bash", |
| "arguments": (tc.get("function") or {}).get("arguments") or "{}", |
| }, |
| } |
| args = tc.get("arguments") |
| if not isinstance(args, str): |
| args = json.dumps(args or {}) |
| return { |
| "id": tc.get("id") or "call_0", |
| "type": "function", |
| "function": {"name": tc.get("name") or "bash", "arguments": args}, |
| } |
|
|
|
|
| def has_edit(msgs) -> bool: |
| for m in msgs: |
| if m.get("role") != "assistant": |
| continue |
| for tc in m.get("tool_calls") or []: |
| fn = (tc.get("function") or {}).get("name") or tc.get("name") |
| if fn in {"edit", "write"}: |
| return True |
| args = (tc.get("function") or {}).get("arguments") or tc.get("arguments") or "" |
| if isinstance(args, dict): |
| args = json.dumps(args) |
| if any(tok in str(args) for tok in ("sed -i", "tee ", "cat >", "git apply", "printf ")): |
| return True |
| return False |
|
|
|
|
| def to_msgs(tr) -> list[dict] | None: |
| nodes = tr.get("nodes") or [] |
| if not (16 <= len(nodes) <= 200): |
| return None |
| msgs = [] |
| for n in nodes: |
| m = n.get("message") or {} |
| role = m.get("role") |
| if role not in {"system", "user", "assistant", "tool"}: |
| continue |
| out = {"role": role} |
| content = m.get("content") |
| if isinstance(content, list): |
| parts = [] |
| for p in content: |
| if isinstance(p, dict) and p.get("type") == "text": |
| parts.append(p.get("text") or "") |
| elif isinstance(p, str): |
| parts.append(p) |
| content = "\n".join(parts) |
| if content is None: |
| content = "" |
| if not isinstance(content, str): |
| content = str(content) |
| out["content"] = content |
| if role == "assistant" and m.get("tool_calls"): |
| out["tool_calls"] = [norm_tc(tc) for tc in m["tool_calls"]] |
| if role == "tool": |
| if m.get("tool_call_id"): |
| out["tool_call_id"] = m["tool_call_id"] |
| elif m.get("name"): |
| out["name"] = m["name"] |
| msgs.append(out) |
| if len(msgs) < 8: |
| return None |
| return msgs |
|
|
|
|
| def main(): |
| rows = [] |
| seen = set() |
| for root in ROOTS: |
| if not root.exists(): |
| continue |
| for f in root.rglob("traces.jsonl"): |
| for line in f.open(): |
| rec = json.loads(line) |
| traces = rec.get("traces") or [rec] |
| for tr in traces: |
| if not isinstance(tr, dict): |
| continue |
| if score(tr, "solved") < 1: |
| continue |
| msgs = to_msgs(tr) |
| if not msgs: |
| continue |
| key = json.dumps([m.get("content", "")[:80] for m in msgs[:6]], sort_keys=True) |
| if key in seen: |
| continue |
| seen.add(key) |
| rows.append({"messages": msgs, "tools": TOOLS, "source": "r2e-solved"}) |
| print("rows", len(rows)) |
| OUT.mkdir(parents=True, exist_ok=True) |
| Dataset.from_list(rows).to_parquet(str(OUT / "train.parquet")) |
| print("wrote", OUT / "train.parquet") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
