dataset_to_sources.py (1069B)
1 import argparse 2 import json 3 import pathlib 4 5 import datasets 6 7 8 def parse_command_argument() -> argparse.Namespace: 9 """ Parse command line arguments. """ 10 parser = argparse.ArgumentParser(description="Extract a jsonl with all sources from a huggingface dataset.") 11 parser.add_argument("output", type=pathlib.Path, help="Path to the output jsonl file.") 12 parser.add_argument("input", type=pathlib.Path, help="Path to the input dataset splits.") 13 14 args: argparse.Namespace = parser.parse_args() 15 return args 16 17 18 def main(output_path: pathlib.Path, input_path: pathlib.Path): 19 with output_path.open("w") as output_file: 20 for split_path in input_path.iterdir(): 21 split = datasets.load_from_disk(str(split_path)) 22 for sample in split: 23 print(json.dumps({"id": sample["id"], "article": sample["source_article"], "headline": sample["source_headline"]}), file=output_file) 24 25 26 if __name__ == "__main__": 27 args: argparse.Namespace = parse_command_argument() 28 main(output_path=args.output, input_path=args.input)