| import argparse |
| import json |
| import random |
| from pathlib import Path |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--questions", default="data/train/verified-mcq.jsonl") |
| parser.add_argument("--exclude", default="data/train/star2/overlap/excluded_ids.txt", |
| help="ids that review_overlap.py found similar to a dev or test question") |
| parser.add_argument("--out_dir", default="data/train/star2") |
| parser.add_argument("--shards", type=int, default=2) |
| parser.add_argument("--seed", type=int, default=0) |
| args = parser.parse_args() |
|
|
| lines = Path(args.questions).read_text(encoding="utf-8").splitlines() |
| excluded = set(Path(args.exclude).read_text(encoding="utf-8").split()) |
| kept = [line for line in lines if json.loads(line)["id"] not in excluded] |
| assert len(lines) - len(kept) == len(excluded), "an excluded id is not in the question file" |
|
|
| random.Random(args.seed).shuffle(kept) |
| out_dir = Path(args.out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| for n in range(args.shards): |
| shard = kept[n :: args.shards] |
| (out_dir / f"shard-{n}.jsonl").write_text("".join(line + "\n" for line in shard), encoding="utf-8") |
| print(f"shard-{n}.jsonl: {len(shard)} questions") |
| print(f"Kept {len(kept)} of {len(lines)} questions; excluded {len(excluded)} listed in {args.exclude}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
