Spaces:
Running
Running
| from __future__ import annotations | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import random | |
| import re | |
| from typing import List | |
| def read_records(path: str) -> List[str]: | |
| text = Path(path).read_text(encoding="utf-8", errors="ignore") | |
| if "<|eos|>" in text: | |
| parts = text.split("<|eos|>") | |
| elif "<|record_end|>" in text: | |
| parts = text.split("<|record_end|>") | |
| elif "<|end|>" in text: | |
| # Backward-compatible fallback for older corpora. New chat-style corpora | |
| # should use <|eos|> as the record separator because <|end|> is also | |
| # used inside conversations to mark turn boundaries. | |
| parts = text.split("<|end|>") | |
| else: | |
| # Fallback for plain corpora: split on blank lines. | |
| parts = re.split(r"\n\s*\n", text) | |
| records = [] | |
| for part in parts: | |
| part = part.strip() | |
| if part: | |
| records.append(part) | |
| return records | |
| def write_records(path: str, records: List[str]) -> None: | |
| out = Path(path) | |
| out.parent.mkdir(parents=True, exist_ok=True) | |
| with out.open("w", encoding="utf-8") as f: | |
| for rec in records: | |
| f.write(rec.strip()) | |
| f.write("\n<|eos|>\n") | |
| def split_corpus( | |
| input_path: str, | |
| train_output: str, | |
| val_output: str, | |
| val_ratio: float = 0.02, | |
| min_val_records: int = 1, | |
| seed: int = 42, | |
| ) -> dict: | |
| if not 0.0 < val_ratio < 0.5: | |
| raise ValueError("val_ratio must be between 0 and 0.5") | |
| records = read_records(input_path) | |
| if len(records) < 2: | |
| raise ValueError("Need at least two records to create a validation split") | |
| rng = random.Random(seed) | |
| rng.shuffle(records) | |
| val_n = max(min_val_records, int(round(len(records) * val_ratio))) | |
| val_n = min(val_n, len(records) - 1) | |
| val_records = records[:val_n] | |
| train_records = records[val_n:] | |
| write_records(train_output, train_records) | |
| write_records(val_output, val_records) | |
| return { | |
| "input": input_path, | |
| "train_output": train_output, | |
| "val_output": val_output, | |
| "records_total": len(records), | |
| "records_train": len(train_records), | |
| "records_val": len(val_records), | |
| "val_ratio_actual": len(val_records) / len(records), | |
| "seed": seed, | |
| } | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Split an Ares corpus into train/validation files.") | |
| parser.add_argument("--input", required=True) | |
| parser.add_argument("--train-output", required=True) | |
| parser.add_argument("--val-output", required=True) | |
| parser.add_argument("--val-ratio", type=float, default=0.02) | |
| parser.add_argument("--min-val-records", type=int, default=1) | |
| parser.add_argument("--seed", type=int, default=42) | |
| args = parser.parse_args() | |
| result = split_corpus( | |
| input_path=args.input, | |
| train_output=args.train_output, | |
| val_output=args.val_output, | |
| val_ratio=args.val_ratio, | |
| min_val_records=args.min_val_records, | |
| seed=args.seed, | |
| ) | |
| print(json.dumps(result, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |