File size: 2,917 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 | #!/usr/bin/env python3
"""Build a deterministic, balanced terminal/SWE SFT mixture."""
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 normalized_component(path: Path, label: str) -> Dataset:
dataset = load_dataset(str(path), split="train").select_columns(COMMON_COLUMNS)
return dataset.add_column("mix_component", [label] * len(dataset))
def write_shards(dataset: Dataset, output_dir: Path, shards: int) -> None:
output_dir.mkdir(parents=True, exist_ok=True)
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("--adapters", type=Path)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--swe-repeats", type=int, default=2)
parser.add_argument("--adapter-repeats", type=int, default=1)
parser.add_argument("--seed", type=int, default=20260814)
parser.add_argument("--shards", type=int, default=64)
args = parser.parse_args()
terminal = normalized_component(args.terminal, "nemotron-terminal-skill")
swe = normalized_component(args.swe, "swe-smith")
parts = [terminal, *([swe] * args.swe_repeats)]
adapters = None
if args.adapters is not None:
adapters = normalized_component(args.adapters, "nemotron-terminal-adapters")
parts.extend([adapters] * args.adapter_repeats)
mixed = concatenate_datasets(parts).shuffle(seed=args.seed)
write_shards(mixed, args.output, min(args.shards, len(mixed)))
report = {
"terminal_path": str(args.terminal),
"terminal_rows": len(terminal),
"swe_path": str(args.swe),
"swe_rows": len(swe),
"swe_repeats": args.swe_repeats,
"adapter_path": str(args.adapters) if args.adapters is not None else None,
"adapter_rows": len(adapters) if adapters is not None else 0,
"adapter_repeats": args.adapter_repeats if adapters is not None else 0,
"output_rows": len(mixed),
"shuffle_seed": args.seed,
"mix": {
"nemotron-terminal-skill": len(terminal),
"swe-smith": len(swe) * args.swe_repeats,
**(
{"nemotron-terminal-adapters": len(adapters) * args.adapter_repeats}
if adapters is not None
else {}
),
},
}
(args.output / "provenance.json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|