| |
| """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() |
|
|
