From 802706dec9d3dad5cf0e7ab5d7771caf845297b9 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 01:08:51 +0000 Subject: [PATCH 1/8] feat: convert episode traces into HF SFT datasets Replace scripts/export_sft.py, which read the pre-episode raw-Trace format, with tools/sft/vf_to_hf.py. The tool reads one traces.jsonl of episode records (written by train runs via FileMonitor and by uv run eval) and emits the dataset shape the SFT trainer consumes: a messages column plus a JSON-encoded tools column. Only trainable agents' traces convert, so a judge's trace never becomes training data. Each branch becomes one sample: a compacted rollout contributes one linear history per branch. Output targets (name, subset, split): locally the tool writes //.parquet and maintains the dataset card's configs metadata so load_dataset(name, subset, split) resolves it; --push pushes to the HF Hub instead. Co-Authored-By: Claude Fable 5 --- scripts/export_sft.py | 94 ----------------------------- tools/sft/vf_to_hf.py | 134 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 134 insertions(+), 94 deletions(-) delete mode 100644 scripts/export_sft.py create mode 100644 tools/sft/vf_to_hf.py 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/sft/vf_to_hf.py b/tools/sft/vf_to_hf.py new file mode 100644 index 0000000000..3bcc90ecf5 --- /dev/null +++ b/tools/sft/vf_to_hf.py @@ -0,0 +1,134 @@ +"""Convert a run's episodes (traces.jsonl) into an HF SFT dataset for `uv run sft`. + +Train and eval runs save episodes as one JSON record per line (prime-rl's FileMonitor +under `rollouts/step_N///traces.jsonl`, `uv run eval` under +`/traces.jsonl`). An episode holds every agent's trace; only trainable agents' +traces become training data (a judge's trace never does). One row per branch: a linear +rollout contributes one sample, a compacted or branched rollout one per branch — each +branch is a linear root-to-leaf history the trainer can feed. + +Rows carry the dataset shape `prime_rl.trainer.sft.data` consumes directly: a `messages` +column (OpenAI chat wire shape) plus a `tools` column (the tools the model was shown, +JSON-encoded — heterogeneous JSON-schema dicts don't fit a fixed Arrow schema). + +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 tools/sft/vf_to_hf.py --name + [--subset default] [--split train] [--min-reward 1.0] [--drop-truncated] [--push] + +Without `--push`, writes `//.parquet` and registers it in the +dataset card (`/README.md`), so the trainer loads it with `--data.name +--data.subsets --data.splits `. Re-running with another subset/split +adds to the same dataset. With `--push`, pushes to the HF Hub repo `` under +config `` and split `` instead. +""" + +import argparse +import json +from pathlib import Path + +import yaml +from datasets import Dataset +from verifiers.v1 import Trace, WireEpisode +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 not trace.agent.trainable: + return False + 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 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() + if text.startswith("---"): + _, header, body = text.split("---", 2) + meta = yaml.safe_load(header) or {} + 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)}---{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="output dataset dir, or HF repo id with --push") + parser.add_argument("--subset", default="default", help="dataset config name") + parser.add_argument("--split", default="train", help="dataset split name") + 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("--push", action="store_true", help="push to the HF Hub instead of writing parquet") + args = parser.parse_args() + + 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: + if keep(trace, args.min_reward, args.drop_truncated): + num_traces += 1 + rows.extend(sft_rows(trace)) + print(f"vf-to-hf: {num_episodes} episode(s) -> {num_traces} trainable trace(s) -> {len(rows)} sample(s)") + if not rows: + raise SystemExit("vf-to-hf: no samples after selection") + + dataset = Dataset.from_list(rows) + if args.push: + dataset.push_to_hub(args.name, config_name=args.subset, split=args.split) + print(f"vf-to-hf: pushed to {args.name} (subset={args.subset}, split={args.split})") + 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"vf-to-hf: wrote {path} -> train with " + f"--data.name {root} --data.subsets {args.subset} --data.splits {args.split}" + ) + + +if __name__ == "__main__": + main() From b93811e1e4a01c3074230e5a9e791fcf6e8d0cf0 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 21:00:29 +0000 Subject: [PATCH 2/8] simplify trace dataset conversion Export every agent branch without selection. Persist trace metadata as\ncolumns so downstream scripts own filtering.\n\nRename the tool to traces_to_hf.py and use traces-to-hf in its\noutput. --- tools/sft/traces_to_hf.py | 180 ++++++++++++++++++++++++++++++++++++++ tools/sft/vf_to_hf.py | 134 ---------------------------- 2 files changed, 180 insertions(+), 134 deletions(-) create mode 100644 tools/sft/traces_to_hf.py delete mode 100644 tools/sft/vf_to_hf.py diff --git a/tools/sft/traces_to_hf.py b/tools/sft/traces_to_hf.py new file mode 100644 index 0000000000..efaefc45c4 --- /dev/null +++ b/tools/sft/traces_to_hf.py @@ -0,0 +1,180 @@ +"""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/sft/traces_to_hf.py --name + [--subset default] [--split train] [--push] + +Without `--push`, the tool writes `//.parquet` and registers it in +`/README.md`. With `--push`, it pushes to the Hugging Face Hub repo `` under +config `` and split ``. +""" + +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), separators=(",", ":")) + + +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() + if text.startswith("---"): + _, header, body = text.split("---", 2) + meta = yaml.safe_load(header) or {} + 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)}---{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="output dataset dir, or HF repo id with --push") + parser.add_argument("--subset", default="default", help="dataset config name") + parser.add_argument("--split", default="train", help="dataset split name") + parser.add_argument("--push", action="store_true", help="push to the HF Hub instead of writing parquet") + args = parser.parse_args() + + 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 args.push: + dataset.push_to_hub(args.name, config_name=args.subset, split=args.split) + print(f"traces-to-hf: pushed to {args.name} (subset={args.subset}, split={args.split})") + 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() diff --git a/tools/sft/vf_to_hf.py b/tools/sft/vf_to_hf.py deleted file mode 100644 index 3bcc90ecf5..0000000000 --- a/tools/sft/vf_to_hf.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Convert a run's episodes (traces.jsonl) into an HF SFT dataset for `uv run sft`. - -Train and eval runs save episodes as one JSON record per line (prime-rl's FileMonitor -under `rollouts/step_N///traces.jsonl`, `uv run eval` under -`/traces.jsonl`). An episode holds every agent's trace; only trainable agents' -traces become training data (a judge's trace never does). One row per branch: a linear -rollout contributes one sample, a compacted or branched rollout one per branch — each -branch is a linear root-to-leaf history the trainer can feed. - -Rows carry the dataset shape `prime_rl.trainer.sft.data` consumes directly: a `messages` -column (OpenAI chat wire shape) plus a `tools` column (the tools the model was shown, -JSON-encoded — heterogeneous JSON-schema dicts don't fit a fixed Arrow schema). - -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 tools/sft/vf_to_hf.py --name - [--subset default] [--split train] [--min-reward 1.0] [--drop-truncated] [--push] - -Without `--push`, writes `//.parquet` and registers it in the -dataset card (`/README.md`), so the trainer loads it with `--data.name ---data.subsets --data.splits `. Re-running with another subset/split -adds to the same dataset. With `--push`, pushes to the HF Hub repo `` under -config `` and split `` instead. -""" - -import argparse -import json -from pathlib import Path - -import yaml -from datasets import Dataset -from verifiers.v1 import Trace, WireEpisode -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 not trace.agent.trainable: - return False - 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 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() - if text.startswith("---"): - _, header, body = text.split("---", 2) - meta = yaml.safe_load(header) or {} - 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)}---{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="output dataset dir, or HF repo id with --push") - parser.add_argument("--subset", default="default", help="dataset config name") - parser.add_argument("--split", default="train", help="dataset split name") - 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("--push", action="store_true", help="push to the HF Hub instead of writing parquet") - args = parser.parse_args() - - 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: - if keep(trace, args.min_reward, args.drop_truncated): - num_traces += 1 - rows.extend(sft_rows(trace)) - print(f"vf-to-hf: {num_episodes} episode(s) -> {num_traces} trainable trace(s) -> {len(rows)} sample(s)") - if not rows: - raise SystemExit("vf-to-hf: no samples after selection") - - dataset = Dataset.from_list(rows) - if args.push: - dataset.push_to_hub(args.name, config_name=args.subset, split=args.split) - print(f"vf-to-hf: pushed to {args.name} (subset={args.subset}, split={args.split})") - 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"vf-to-hf: wrote {path} -> train with " - f"--data.name {root} --data.subsets {args.subset} --data.splits {args.split}" - ) - - -if __name__ == "__main__": - main() From fb01523cffa37cfc16a0e29cb0b7dba805c92569 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 21:03:00 +0000 Subject: [PATCH 3/8] move trace converter under tools --- tools/{sft/traces_to_hf.py => convert_traces_to_hf_dataset.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename tools/{sft/traces_to_hf.py => convert_traces_to_hf_dataset.py} (98%) diff --git a/tools/sft/traces_to_hf.py b/tools/convert_traces_to_hf_dataset.py similarity index 98% rename from tools/sft/traces_to_hf.py rename to tools/convert_traces_to_hf_dataset.py index efaefc45c4..ae67ffdd6e 100644 --- a/tools/sft/traces_to_hf.py +++ b/tools/convert_traces_to_hf_dataset.py @@ -9,7 +9,7 @@ so later scripts can filter them directly. Metadata with variable schemas stays as JSON. Usage (from the prime-rl repo): - uv run python tools/sft/traces_to_hf.py --name + uv run python tools/convert_traces_to_hf_dataset.py --name [--subset default] [--split train] [--push] Without `--push`, the tool writes `//.parquet` and registers it in From 82c1d648e0e4b6f1d1765bb0ffaf8ba78c050984 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 21:10:29 +0000 Subject: [PATCH 4/8] push converted datasets by default Treat the dataset name as a Hub repository unless --local is set.\nLet --private create a private Hub repository. --- tools/convert_traces_to_hf_dataset.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/tools/convert_traces_to_hf_dataset.py b/tools/convert_traces_to_hf_dataset.py index ae67ffdd6e..f476b9ed0f 100644 --- a/tools/convert_traces_to_hf_dataset.py +++ b/tools/convert_traces_to_hf_dataset.py @@ -10,11 +10,11 @@ Usage (from the prime-rl repo): uv run python tools/convert_traces_to_hf_dataset.py --name - [--subset default] [--split train] [--push] + [--subset default] [--split train] [--private] [--local] -Without `--push`, the tool writes `//.parquet` and registers it in -`/README.md`. With `--push`, it pushes to the Hugging Face Hub repo `` under -config `` and split ``. +By default, the tool pushes to the Hugging Face Hub repo ``. Use `--private` when +creating a private repo. With `--local`, it writes `//.parquet` and +registers it in `/README.md` instead. """ import argparse @@ -142,11 +142,14 @@ def register_in_dataset_card(root: Path, subset: str, split: str, rel_path: str) 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="output dataset dir, or HF repo id with --push") + 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("--push", action="store_true", help="push to the HF Hub instead of writing parquet") + parser.add_argument("--private", action="store_true", help="make a new Hub dataset private") + parser.add_argument("--local", action="store_true", help="write parquet locally instead of pushing") args = parser.parse_args() + if args.local and args.private: + parser.error("--private cannot be used with --local") num_episodes, num_traces, rows = 0, 0, [] with args.traces.open(encoding="utf-8") as f: @@ -163,9 +166,14 @@ def main() -> None: raise SystemExit("traces-to-hf: no branches found") dataset = Dataset.from_list(rows) - if args.push: - dataset.push_to_hub(args.name, config_name=args.subset, split=args.split) - print(f"traces-to-hf: pushed to {args.name} (subset={args.subset}, split={args.split})") + if not args.local: + dataset.push_to_hub( + args.name, + config_name=args.subset, + split=args.split, + private=True if args.private else None, + ) + print(f"traces-to-hf: pushed to {args.name} (subset={args.subset}, split={args.split}, private={args.private})") return root = Path(args.name) rel_path = f"{args.subset}/{args.split}.parquet" From 08c1be00a1b2233f1109a1f2bcd03ce5d8063397 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 21:13:46 +0000 Subject: [PATCH 5/8] document one-item list overrides --- skills/configs/SKILL.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index a7f0abfd18..9b5edf29d8 100644 --- a/skills/configs/SKILL.md +++ b/skills/configs/SKILL.md @@ -60,6 +60,8 @@ env.agent.runtime.type = "subprocess" ``` CLI: `--orchestrator.train.source.0.env.taskset.id reverse-text` or `--orchestrator.eval.source.0.env.taskset.id reverse-text`. +Use a JSON literal for a one-item list. For example, use `--data.subsets '["default"]'` and +`--data.splits '["train"]'`. A bare scalar does not pass list validation. The `sft` entrypoint takes the same eval shape at the top level for online evals: `[eval]` + `[[eval.source]]` (with `[inference]` for the server), e.g. `--eval.source.0.env.taskset.id reverse-text`. From 8b4ae0aa4e4e33d57dcc105d34f9d83e86b22f11 Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 21:31:17 +0000 Subject: [PATCH 6/8] preserve existing dataset cards Keep README content when it has no valid YAML frontmatter. Continue\nupdating valid metadata without replacing the card body. --- tools/convert_traces_to_hf_dataset.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tools/convert_traces_to_hf_dataset.py b/tools/convert_traces_to_hf_dataset.py index f476b9ed0f..0cdd238bf4 100644 --- a/tools/convert_traces_to_hf_dataset.py +++ b/tools/convert_traces_to_hf_dataset.py @@ -122,9 +122,14 @@ def register_in_dataset_card(root: Path, subset: str, split: str, rel_path: str) body = "" if readme.exists(): text = readme.read_text() - if text.startswith("---"): - _, header, body = text.split("---", 2) + 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: @@ -135,7 +140,7 @@ def register_in_dataset_card(root: Path, subset: str, split: str, rel_path: str) config["data_files"].append({"split": split, "path": rel_path}) else: entry["path"] = rel_path - text = f"---\n{yaml.safe_dump(meta, sort_keys=False)}---{body}" + text = f"---\n{yaml.safe_dump(meta, sort_keys=False)}---\n{body}" readme.write_text(text if text.endswith("\n") else text + "\n") From c479201be335e7bb28618e89c72a36721a923afb Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 22:15:46 +0000 Subject: [PATCH 7/8] address converter review --- skills/configs/SKILL.md | 2 -- tools/convert_traces_to_hf_dataset.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index 9b5edf29d8..a7f0abfd18 100644 --- a/skills/configs/SKILL.md +++ b/skills/configs/SKILL.md @@ -60,8 +60,6 @@ env.agent.runtime.type = "subprocess" ``` CLI: `--orchestrator.train.source.0.env.taskset.id reverse-text` or `--orchestrator.eval.source.0.env.taskset.id reverse-text`. -Use a JSON literal for a one-item list. For example, use `--data.subsets '["default"]'` and -`--data.splits '["train"]'`. A bare scalar does not pass list validation. The `sft` entrypoint takes the same eval shape at the top level for online evals: `[eval]` + `[[eval.source]]` (with `[inference]` for the server), e.g. `--eval.source.0.env.taskset.id reverse-text`. diff --git a/tools/convert_traces_to_hf_dataset.py b/tools/convert_traces_to_hf_dataset.py index 0cdd238bf4..1a5dff3acf 100644 --- a/tools/convert_traces_to_hf_dataset.py +++ b/tools/convert_traces_to_hf_dataset.py @@ -40,7 +40,7 @@ def jsonable(value: Any) -> Any: def as_json(value: Any) -> str: - return json.dumps(jsonable(value), separators=(",", ":")) + return json.dumps(jsonable(value)) def trace_rows(episode: WireEpisode, trace: Trace) -> list[dict]: From 8834b7aea61f2e6e9ba467758b593cf6ffb7a29b Mon Sep 17 00:00:00 2001 From: Mika Senghaas Date: Tue, 25 Aug 2026 23:08:53 +0000 Subject: [PATCH 8/8] make hub datasets private by default --- tools/convert_traces_to_hf_dataset.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tools/convert_traces_to_hf_dataset.py b/tools/convert_traces_to_hf_dataset.py index 1a5dff3acf..35be68750c 100644 --- a/tools/convert_traces_to_hf_dataset.py +++ b/tools/convert_traces_to_hf_dataset.py @@ -10,10 +10,10 @@ Usage (from the prime-rl repo): uv run python tools/convert_traces_to_hf_dataset.py --name - [--subset default] [--split train] [--private] [--local] + [--subset default] [--split train] [--public] [--local] -By default, the tool pushes to the Hugging Face Hub repo ``. Use `--private` when -creating a private repo. With `--local`, it writes `//.parquet` and +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. """ @@ -150,11 +150,11 @@ def main() -> None: 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("--private", action="store_true", help="make a new Hub dataset private") + 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.private: - parser.error("--private cannot be used with --local") + 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: @@ -176,9 +176,11 @@ def main() -> None: args.name, config_name=args.subset, split=args.split, - private=True if args.private else None, + private=not args.public, + ) + print( + f"traces-to-hf: pushed to {args.name} (subset={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={args.private})") return root = Path(args.name) rel_path = f"{args.subset}/{args.split}.parquet"