File size: 3,618 Bytes
9589849 | 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 | #!/usr/bin/env python3
"""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()
|