wehub-resource-sync
128 lines
4.0 KiB
Python
128 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Merge the first 5 filtered-HN chunks and split them into train/eval/test JSONL."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
|
|
|
|
DEFAULT_CHUNK_INPUTS = [
|
|
Path(
|
|
"/home/user/wiki-screenshot-training/training/data/lite-query-v2-full-filtered-hn-v2-chunks/chunk_000000_009999/filtered_hn.jsonl"
|
|
),
|
|
Path(
|
|
"/home/user/wiki-screenshot-training/training/data/lite-query-v2-full-filtered-hn-v2-chunks/chunk_010000_019999/filtered_hn.jsonl"
|
|
),
|
|
Path(
|
|
"/home/user/wiki-screenshot-training/training/data/lite-query-v2-full-filtered-hn-v2-chunks/chunk_020000_029999/filtered_hn.jsonl"
|
|
),
|
|
Path(
|
|
"/home/user/wiki-screenshot-training/training/data/lite-query-v2-full-filtered-hn-v2-chunks/chunk_030000_039999/filtered_hn.jsonl"
|
|
),
|
|
Path(
|
|
"/home/user/wiki-screenshot-training/training/data/lite-query-v2-full-filtered-hn-v2-chunks/chunk_040000_049999/filtered_hn.jsonl"
|
|
),
|
|
]
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--inputs",
|
|
nargs="+",
|
|
default=[str(path) for path in DEFAULT_CHUNK_INPUTS],
|
|
help="Filtered chunk JSONL files to merge before splitting.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
default="/home/user/wiki-screenshot-training/training/data/lite-query-v2-full-filtered-hn-v2-chunks/split",
|
|
help="Where to write train/eval/test JSONL files.",
|
|
)
|
|
parser.add_argument("--train-ratio", type=float, default=0.90)
|
|
parser.add_argument("--eval-ratio", type=float, default=0.05)
|
|
parser.add_argument("--test-ratio", type=float, default=0.05)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
return parser.parse_args()
|
|
|
|
|
|
def read_jsonl(path: Path) -> list[dict]:
|
|
rows = []
|
|
with path.open() as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
rows.append(json.loads(line))
|
|
return rows
|
|
|
|
|
|
def write_jsonl(path: Path, rows: list[dict]) -> None:
|
|
with path.open("w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
input_paths = [Path(path) for path in args.inputs]
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
ratio_sum = args.train_ratio + args.eval_ratio + args.test_ratio
|
|
if abs(ratio_sum - 1.0) > 1e-9:
|
|
raise ValueError(f"Ratios must sum to 1.0, got {ratio_sum}")
|
|
|
|
rows = []
|
|
input_counts = {}
|
|
for path in input_paths:
|
|
chunk_rows = read_jsonl(path)
|
|
rows.extend(chunk_rows)
|
|
input_counts[str(path)] = len(chunk_rows)
|
|
|
|
rng = random.Random(args.seed)
|
|
rng.shuffle(rows)
|
|
|
|
total = len(rows)
|
|
train_end = int(total * args.train_ratio)
|
|
eval_end = train_end + int(total * args.eval_ratio)
|
|
|
|
train_rows = rows[:train_end]
|
|
eval_rows = rows[train_end:eval_end]
|
|
test_rows = rows[eval_end:]
|
|
|
|
train_path = output_dir / "train_hn.jsonl"
|
|
eval_path = output_dir / "eval_hn.jsonl"
|
|
test_path = output_dir / "test_hn.jsonl"
|
|
summary_path = output_dir / "split_summary.json"
|
|
|
|
write_jsonl(train_path, train_rows)
|
|
write_jsonl(eval_path, eval_rows)
|
|
write_jsonl(test_path, test_rows)
|
|
|
|
summary = {
|
|
"inputs": input_counts,
|
|
"seed": args.seed,
|
|
"total_rows": total,
|
|
"train_ratio": args.train_ratio,
|
|
"eval_ratio": args.eval_ratio,
|
|
"test_ratio": args.test_ratio,
|
|
"train_rows": len(train_rows),
|
|
"eval_rows": len(eval_rows),
|
|
"test_rows": len(test_rows),
|
|
"output_files": {
|
|
"train": str(train_path),
|
|
"eval": str(eval_path),
|
|
"test": str(test_path),
|
|
},
|
|
}
|
|
summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True))
|
|
|
|
print(json.dumps(summary, indent=2, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|