diff --git a/scripts/export_sft.py b/scripts/export_sft.py deleted file mode 100644 index ece9bf0b80..0000000000 --- a/scripts/export_sft.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Export a verifiers v1 eval run (traces.jsonl) as an SFT dataset for `uv run sft`. - -Reads a finished run's saved traces and reshapes them into the dataset shape the SFT trainer -consumes directly (see `prime_rl.trainer.sft.data`): a `messages` column (OpenAI chat wire -shape) plus a `tools` column (the tools the model was shown, from `Trace.tools`, -JSON-encoded — heterogeneous JSON-schema dicts don't fit a fixed Arrow schema). One row per -branch: a linear rollout contributes one sample, a compacted/subagent rollout one per branch -(one training sample is built per branch). - -Selection: generation-errored traces (`stop_condition == "error"`) always drop — a broken -transcript is not a sample. A scoring-only error keeps the generation outcome as its stop -condition and a complete conversation, so it stays; its reward may be partial/zero, which -`--min-reward` handles. - -Usage (from the prime-rl repo): - uv run python scripts/export_sft.py [--min-reward 1.0] [--drop-truncated] - [-o OUT_DIR] [--push HF_REPO_ID] - -Writes `/sft/train.parquet` by default — point the trainer at it with -`--data.name /sft`. Requires a verifiers release carrying `Trace.tools` -(PrimeIntellect-ai/verifiers#1963). -""" - -import argparse -import json -from pathlib import Path - -from datasets import Dataset -from verifiers.v1 import Trace, WireTrace -from verifiers.v1.dialects.chat import message_to_wire - - -def sft_rows(trace: Trace) -> list[dict]: - """A trace's SFT rows — one per branch: the branch's conversation as OpenAI chat wire - dicts plus the trace's advertised tools, JSON-encoded.""" - tools = json.dumps([t.model_dump(mode="json", exclude_none=True) for t in trace.tools or []]) - return [ - { - "messages": [message_to_wire(m) for m in branch.messages], - "tools": tools, - } - for branch in trace.branches - if branch.messages - ] - - -def keep(trace: Trace, min_reward: float | None, drop_truncated: bool) -> bool: - """Whether a trace is worth training on (see module docstring for the error semantics).""" - if trace.stop_condition == "error": - return False - if drop_truncated and trace.is_truncated: - return False - return min_reward is None or trace.reward >= min_reward - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - parser.add_argument("run_dir", type=Path, help="the eval run dir (holds traces.jsonl)") - parser.add_argument("--min-reward", type=float, default=None, help="keep traces with reward >= this") - parser.add_argument("--drop-truncated", action="store_true", help="drop budget-cut traces") - parser.add_argument("-o", "--output-dir", type=Path, default=None, help="default: /sft") - parser.add_argument("--push", default=None, help="HF repo id to push to instead of writing parquet") - args = parser.parse_args() - - traces_path = args.run_dir / "traces.jsonl" - if not traces_path.exists(): - raise SystemExit(f"no traces.jsonl in {args.run_dir}") - - total, rows = 0, [] - with traces_path.open(encoding="utf-8") as f: - for line in f: - if not line.strip(): - continue - total += 1 - trace = WireTrace.model_validate(json.loads(line)) - if keep(trace, args.min_reward, args.drop_truncated): - rows.extend(sft_rows(trace)) - print(f"export-sft: {total} trace(s) -> {len(rows)} row(s)") - if not rows: - raise SystemExit("export-sft: no rows to export after selection") - - dataset = Dataset.from_list(rows) - if args.push: - dataset.push_to_hub(args.push) - print(f"export-sft: pushed to {args.push}") - return - out = args.output_dir or args.run_dir / "sft" - out.mkdir(parents=True, exist_ok=True) - dataset.to_parquet(str(out / "train.parquet")) - print(f"export-sft: wrote {out / 'train.parquet'} -> train with --data.name {out}") - - -if __name__ == "__main__": - main() diff --git a/tools/convert_traces_to_hf_dataset.py b/tools/convert_traces_to_hf_dataset.py new file mode 100644 index 0000000000..35be68750c --- /dev/null +++ b/tools/convert_traces_to_hf_dataset.py @@ -0,0 +1,195 @@ +"""Convert episode traces into a branch-level Hugging Face dataset. + +Train and eval runs save one episode per line in `traces.jsonl`. This tool writes one +dataset row for every branch of every agent trace. It does not select or filter rows. + +Each row includes `messages` and JSON-encoded `tools` for SFT. It also includes trace, +agent, task, run, outcome, error, timing, and usage metadata. Scalar outcome fields such +as `reward`, `stop_condition`, `has_error`, and `is_truncated` stay as top-level columns +so later scripts can filter them directly. Metadata with variable schemas stays as JSON. + +Usage (from the prime-rl repo): + uv run python tools/convert_traces_to_hf_dataset.py --name + [--subset default] [--split train] [--public] [--local] + +By default, the tool creates a private Hugging Face Hub repo named ``. Use `--public` +to create a public repo. With `--local`, it writes `//.parquet` and +registers it in `/README.md` instead. +""" + +import argparse +import json +from pathlib import Path +from typing import Any + +import yaml +from datasets import Dataset +from pydantic import BaseModel +from verifiers.v1 import Trace, WireEpisode +from verifiers.v1.dialects.chat import message_to_wire + + +def jsonable(value: Any) -> Any: + if isinstance(value, BaseModel): + return value.model_dump(mode="json", exclude_none=True) + if isinstance(value, list): + return [jsonable(item) for item in value] + if isinstance(value, dict): + return {key: jsonable(item) for key, item in value.items()} + return value + + +def as_json(value: Any) -> str: + return json.dumps(jsonable(value)) + + +def trace_rows(episode: WireEpisode, trace: Trace) -> list[dict]: + """Convert each branch of one trace without selecting or filtering it.""" + last_error = trace.last_error + run = episode.run + work = getattr(run, "work", None) + policy = getattr(work, "policy", None) + common = { + "episode_id": episode.id, + "episode_ok": episode.ok, + "episode_has_error": bool(episode.errors), + "episode_errors": as_json(episode.errors), + "env_id": episode.env.id, + "env_name": episode.env.name, + "group_id": episode.group.id if episode.group else None, + "run_type": run.type if run else None, + "run_id": run.id if run else None, + "run_name": run.name if run else None, + "work_type": work.type if work else None, + "step": getattr(work, "step", None), + "policy_start": policy.start if policy else None, + "policy_end": policy.end if policy else None, + "trace_id": trace.id, + "trace_version": trace.version, + "verifiers_version": trace.verifiers.version, + "verifiers_commit": trace.verifiers.commit, + "task_type": trace.task.type, + "task_data": as_json(trace.task.data), + "task_key": trace.task.key, + "task_hash": trace.task.hash, + "agent": trace.agent.name, + "trainable": trace.agent.trainable, + "agent_config": as_json(trace.agent.config), + "agent_runtime": as_json(trace.agent.runtime), + "tools": as_json(trace.tools), + "reward": trace.reward, + "rewards": as_json(trace.rewards), + "metrics": as_json(trace.metrics), + "info": as_json(trace.info), + "is_completed": trace.is_completed, + "ok": trace.ok, + "stop_condition": trace.stop_condition, + "has_error": trace.has_error, + "error_type": last_error.type if last_error else None, + "error_message": last_error.message if last_error else None, + "error_status_code": last_error.status_code if last_error else None, + "errors": as_json(trace.errors), + "is_truncated": trace.is_truncated, + "request_rewrites": as_json(trace.request_rewrites), + "response_rewrites": as_json(trace.response_rewrites), + "extra_usage": as_json(trace.extra_usage), + "timing": as_json(trace.timing), + "num_branches": trace.num_branches, + "num_turns": trace.num_turns, + "num_input_tokens": trace.num_input_tokens, + "num_output_tokens": trace.num_output_tokens, + "num_total_tokens": trace.num_total_tokens, + } + return [ + { + **common, + "branch_index": branch.index, + "messages": [message_to_wire(m) for m in branch.messages], + "calls": as_json(branch.calls), + "branch_num_input_tokens": branch.num_input_tokens, + "branch_num_output_tokens": branch.num_output_tokens, + "branch_num_total_tokens": branch.num_total_tokens, + } + for branch in trace.branches + ] + + +def register_in_dataset_card(root: Path, subset: str, split: str, rel_path: str) -> None: + """Point the dataset card's `configs` metadata at the parquet, so + `load_dataset(root, subset, split=split)` resolves it.""" + readme = root / "README.md" + meta: dict = {} + body = "" + if readme.exists(): + text = readme.read_text() + lines = text.splitlines(keepends=True) + end = next((i for i, line in enumerate(lines[1:], 1) if line.strip() == "---"), None) + if lines and lines[0].strip() == "---" and end is not None: + header = "".join(lines[1:end]) + body = "".join(lines[end + 1 :]) + meta = yaml.safe_load(header) or {} + else: + body = text + configs = meta.setdefault("configs", []) + config = next((c for c in configs if c["config_name"] == subset), None) + if config is None: + config = {"config_name": subset, "data_files": []} + configs.append(config) + entry = next((e for e in config["data_files"] if e["split"] == split), None) + if entry is None: + config["data_files"].append({"split": split, "path": rel_path}) + else: + entry["path"] = rel_path + text = f"---\n{yaml.safe_dump(meta, sort_keys=False)}---\n{body}" + readme.write_text(text if text.endswith("\n") else text + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("traces", type=Path, help="a run's traces.jsonl (one episode per line)") + parser.add_argument("--name", required=True, help="HF repo id, or output dataset dir with --local") + parser.add_argument("--subset", default="default", help="dataset config name") + parser.add_argument("--split", default="train", help="dataset split name") + parser.add_argument("--public", action="store_true", help="make a new Hub dataset public") + parser.add_argument("--local", action="store_true", help="write parquet locally instead of pushing") + args = parser.parse_args() + if args.local and args.public: + parser.error("--public cannot be used with --local") + + num_episodes, num_traces, rows = 0, 0, [] + with args.traces.open(encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + num_episodes += 1 + episode = WireEpisode.model_validate(json.loads(line)) + for trace in episode.traces: + num_traces += 1 + rows.extend(trace_rows(episode, trace)) + print(f"traces-to-hf: {num_episodes} episode(s) -> {num_traces} trace(s) -> {len(rows)} branch(es)") + if not rows: + raise SystemExit("traces-to-hf: no branches found") + + dataset = Dataset.from_list(rows) + if not args.local: + dataset.push_to_hub( + args.name, + config_name=args.subset, + split=args.split, + private=not args.public, + ) + print( + f"traces-to-hf: pushed to {args.name} (subset={args.subset}, split={args.split}, private={not args.public})" + ) + return + root = Path(args.name) + rel_path = f"{args.subset}/{args.split}.parquet" + path = root / rel_path + path.parent.mkdir(parents=True, exist_ok=True) + dataset.to_parquet(str(path)) + register_in_dataset_card(root, args.subset, args.split, rel_path) + print(f"traces-to-hf: wrote {path} (subset={args.subset}, split={args.split})") + + +if __name__ == "__main__": + main()