()
| 48 | |
| 49 | |
| 50 | def main(): |
| 51 | parser = argparse.ArgumentParser() |
| 52 | parser.add_argument("--input", required=True) |
| 53 | parser.add_argument("--output-dir", default="uncommon_route/data/v2_splits") |
| 54 | parser.add_argument("--seed", type=int, default=42) |
| 55 | args = parser.parse_args() |
| 56 | |
| 57 | rows = [] |
| 58 | with open(args.input, encoding="utf-8") as f: |
| 59 | for line in f: |
| 60 | line = line.strip() |
| 61 | if line: |
| 62 | rows.append(json.loads(line)) |
| 63 | |
| 64 | train, cal, holdout = stratified_3way_split(rows, seed=args.seed) |
| 65 | out_dir = Path(args.output_dir) |
| 66 | out_dir.mkdir(parents=True, exist_ok=True) |
| 67 | |
| 68 | for name, split in [("train", train), ("calibration", cal), ("holdout", holdout)]: |
| 69 | path = out_dir / f"{name}.jsonl" |
| 70 | with open(path, "w", encoding="utf-8") as f: |
| 71 | for row in split: |
| 72 | f.write(json.dumps(row, ensure_ascii=False) + "\n") |
| 73 | print(f"{name}: {len(split)} rows -> {path}") |
| 74 | print(f"\nTotal: {len(train)} + {len(cal)} + {len(holdout)} = {len(train) + len(cal) + len(holdout)}") |
| 75 | |
| 76 | |
| 77 | if __name__ == "__main__": |
no test coverage detected