convert_predictions.py (1758B)
1 import argparse 2 import contextlib 3 import itertools 4 import json 5 import pathlib 6 7 import datasets 8 9 10 def parse_command_argument() -> argparse.Namespace: 11 """ Parse command line arguments. """ 12 parser = argparse.ArgumentParser(description="Convert a DEGREE predictions file into an AEC predictions file.") 13 parser.add_argument("output", type=pathlib.Path, help="Path to the output AEC predictions.") 14 parser.add_argument("input", type=pathlib.Path, help="Path to the input DEGREE predictions.") 15 parser.add_argument("split", type=pathlib.Path, help="Path to the AEC split that resulted in the given predictions.") 16 17 args: argparse.Namespace = parser.parse_args() 18 return args 19 20 21 def main(output_path: pathlib.Path, input_path: pathlib.Path, split_path: pathlib.Path) -> None: 22 split: datasets.Dataset = datasets.load_from_disk(str(split_path)) 23 with input_path.open("r") as input_file: 24 degree_predictions = json.load(input_file) 25 26 with output_path.open("w") as output_file: 27 for sample_id, degree_object in zip(split["id"], degree_predictions): 28 arguments: dict[str, str] = {"id": sample_id} 29 for role, argument in degree_object["pred arguments"].items(): 30 argument = argument.strip() 31 if role.startswith("deaths_"): 32 with contextlib.suppress(ValueError): 33 arguments[role] = int("".join(itertools.takewhile(str.isdigit, argument))) 34 else: 35 arguments[role] = argument 36 print(json.dumps(arguments), file=output_file) 37 38 39 if __name__ == "__main__": 40 args: argparse.Namespace = parse_command_argument() 41 main(output_path=args.output, input_path=args.input, split_path=args.split)