UCDP-AEC

Models and data for the UCDP Abstractive Event analysis Corpus paper
git clone https://esimon.eu/repos/UCDP-AEC.git
Log | Files | Refs | README | LICENSE

convert_aec_preds.py (2392B)

      1 from typing import Any
      2 import argparse
      3 import contextlib
      4 import json
      5 import pathlib
      6 
      7 import datasets
      8 
      9 from ucdp_fields import UCDP_FIELDS
     10 
     11 
     12 def parse_command_argument() -> argparse.Namespace:
     13     """ Parse command line arguments. """
     14     parser = argparse.ArgumentParser(description="Convert a Text2Event predictions file into an AEC predictions file.")
     15     parser.add_argument("output", type=pathlib.Path, help="Path to the output AEC predictions.")
     16     parser.add_argument("input", type=pathlib.Path, help="Path to the input Text2Event predictions.")
     17     parser.add_argument("split", type=pathlib.Path, help="Path to the AEC split that resulted in the given predictions.")
     18 
     19     args: argparse.Namespace = parser.parse_args()
     20     return args
     21 
     22 
     23 def main(output_path: pathlib.Path, input_path: pathlib.Path, split_path: pathlib.Path) -> None:
     24     split: datasets.Dataset = datasets.load_from_disk(str(split_path))
     25     with output_path.open("w") as output_file, input_path.open("r") as input_file:
     26         for sample_id, input_line in zip(split["id"], input_file):
     27             input_line = input_line.strip()
     28             if input_line.startswith("<extra_id_0> <extra_id_0> "):
     29                 input_line = input_line[len("<extra_id_0> <extra_id_0> "):]
     30             if input_line.endswith("<extra_id_1> <extra_id_1>"):
     31                 input_line = input_line[:-len("<extra_id_1> <extra_id_1>")]
     32             children: list[str] = input_line.split("<extra_id_1> <extra_id_0> ")
     33             arguments: dict[str, str] = {"id": sample_id}
     34             for child in children:
     35                 role: str | None = None
     36                 for key, label, _ in sorted(UCDP_FIELDS, key=lambda field: len(field[1]), reverse=True):
     37                     if child.startswith(label):
     38                         role = key
     39                         break
     40                 if role is not None:
     41                     argument: str = child[len(label)+1:]
     42                     if role.startswith("deaths_"):
     43                         with contextlib.suppress(ValueError):
     44                             arguments[role] = int(argument)
     45                     else:
     46                         arguments[role] = argument
     47             print(json.dumps(arguments), file=output_file)
     48 
     49 
     50 if __name__ == "__main__":
     51     args: argparse.Namespace = parse_command_argument()
     52     main(output_path=args.output, input_path=args.input, split_path=args.split)