| |
| """Build a deterministic SFT mixture containing only complete trajectories. |
| |
| Inputs are already-decontaminated, public trajectory datasets. Completion is |
| defined structurally and without reference to evaluation tasks: the final |
| message must be a non-tool assistant response. This excludes source rollouts |
| that ended at their turn limit or immediately after a shell observation. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from datasets import Dataset, concatenate_datasets, load_dataset |
|
|
| COMMON_COLUMNS = ["messages", "tools", "source", "num_tool_calls"] |
|
|
|
|
| def is_complete(messages: list[dict]) -> bool: |
| if not messages: |
| return False |
| final = messages[-1] |
| return ( |
| final.get("role") == "assistant" |
| and not final.get("tool_calls") |
| and bool((final.get("content") or "").strip()) |
| ) |
|
|
|
|
| def component(path: Path, label: str, num_proc: int) -> tuple[Dataset, int]: |
| dataset = load_dataset(str(path), split="train").select_columns(COMMON_COLUMNS) |
| input_rows = len(dataset) |
| dataset = dataset.filter( |
| is_complete, |
| input_columns=["messages"], |
| num_proc=num_proc, |
| desc=f"Keeping complete {label} trajectories", |
| ) |
| dataset = dataset.add_column("mix_component", [label] * len(dataset)) |
| return dataset, input_rows |
|
|
|
|
| def write_shards(dataset: Dataset, output_dir: Path, shards: int) -> None: |
| output_dir.mkdir(parents=True, exist_ok=False) |
| for index in range(shards): |
| dataset.shard(shards, index, contiguous=True).to_parquet( |
| output_dir / f"train-{index:05d}-of-{shards:05d}.parquet" |
| ) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--terminal", type=Path, required=True) |
| parser.add_argument("--swe", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--terminal-repeats", type=int, default=2) |
| parser.add_argument("--swe-repeats", type=int, default=2) |
| parser.add_argument("--seed", type=int, default=20260814) |
| parser.add_argument("--shards", type=int, default=64) |
| parser.add_argument("--num-proc", type=int, default=32) |
| args = parser.parse_args() |
|
|
| terminal, terminal_input_rows = component( |
| args.terminal, "nemotron-terminal-skill-complete", args.num_proc |
| ) |
| swe, swe_input_rows = component(args.swe, "swe-smith-complete", args.num_proc) |
| parts = [terminal] * args.terminal_repeats + [swe] * args.swe_repeats |
| mixed = concatenate_datasets(parts).shuffle(seed=args.seed) |
| write_shards(mixed, args.output, min(args.shards, len(mixed))) |
|
|
| report = { |
| "selection": "final message is a non-tool assistant response with non-empty content", |
| "selection_uses_evaluation_content": False, |
| "terminal_path": str(args.terminal), |
| "terminal_input_rows": terminal_input_rows, |
| "terminal_complete_rows": len(terminal), |
| "terminal_repeats": args.terminal_repeats, |
| "swe_path": str(args.swe), |
| "swe_input_rows": swe_input_rows, |
| "swe_complete_rows": len(swe), |
| "swe_repeats": args.swe_repeats, |
| "output_rows": len(mixed), |
| "shuffle_seed": args.seed, |
| "mix": { |
| "nemotron-terminal-skill-complete": len(terminal) * args.terminal_repeats, |
| "swe-smith-complete": len(swe) * args.swe_repeats, |
| }, |
| } |
| (args.output / "provenance.json").write_text(json.dumps(report, indent=2) + "\n") |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
