File size: 1,432 Bytes
044f4e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | 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()
|