build_cache.py (1379B)
1 import argparse 2 import json 3 import pathlib 4 5 import datasets 6 7 from hplt_align.common import map_minhash 8 9 10 def parse_command_argument() -> argparse.Namespace: 11 """ Parse command line arguments. """ 12 parser = argparse.ArgumentParser(description="Update a huggingface dataset using a minhash input and output.") 13 parser.add_argument("output", type=pathlib.Path, help="Path to the cache file.") 14 parser.add_argument("dataset", type=pathlib.Path, help="Path to the target dataset splits.") 15 parser.add_argument("minhash", type=pathlib.Path, help="Path to the output zip of the minhash script.") 16 17 args: argparse.Namespace = parser.parse_args() 18 return args 19 20 21 22 def main(output_path: pathlib.Path, dataset_path: pathlib.Path, minhash_path: pathlib.Path) -> None: 23 minhash: dict[str, str] = {} 24 map_minhash(minhash_path, lambda match: minhash.update({match["qid"]: match["text"]})) 25 26 with output_path.open("w") as output_file: 27 for split_path in dataset_path.iterdir(): 28 split = datasets.load_from_disk(str(split_path)) 29 for hplt_id in split["source_article"]: 30 print(json.dumps({"id": hplt_id, "text": minhash[hplt_id]}), file=output_file) 31 32 33 if __name__ == "__main__": 34 args: argparse.Namespace = parse_command_argument() 35 main(output_path=args.output, dataset_path=args.dataset, minhash_path=args.minhash)