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_dataset.py (2160B)

      1 from typing import Any
      2 import argparse
      3 import json
      4 import pathlib
      5 
      6 import datasets
      7 
      8 from ucdp_fields import UCDP_FIELDS
      9 
     10 
     11 SPLIT_NAME_MAP: dict[str, str] = {
     12     "train": "train",
     13     "validation": "val",
     14     "test": "test"
     15 }
     16 
     17 
     18 def parse_command_argument() -> argparse.Namespace:
     19     """ Parse command line arguments. """
     20     parser = argparse.ArgumentParser(description="Convert a AEC dataset into a Text2Event dataset.")
     21     parser.add_argument("output", type=pathlib.Path, help="Path to the output Text2Event dataset directory.")
     22     parser.add_argument("input", type=pathlib.Path, help="Path to the input AEC dataset.")
     23 
     24     args: argparse.Namespace = parser.parse_args()
     25     return args
     26 
     27 
     28 def write_schema(output_path: pathlib.Path) -> None:
     29     with output_path.open("w") as output_file:
     30         print(json.dumps(["UCDP"]), file=output_file)
     31         print(json.dumps([label for _, label, _ in UCDP_FIELDS]), file=output_file)
     32         print(json.dumps({"UCDP": [label for _, label, _ in UCDP_FIELDS]}), file=output_file)
     33 
     34 
     35 def input_template(sample: dict[str, Any]) -> str:
     36     return f"{sample['source_date']} {sample['source_article']}"
     37 
     38 
     39 def output_template(sample: dict[str, Any]) -> str:
     40     return "<extra_id_0> " + " ".join(f"<extra_id_0> {label} {sample[field]} <extra_id_1>" for field, label, _ in UCDP_FIELDS) + " <extra_id_1>"
     41 
     42 
     43 def convert_split(output_path: pathlib.Path, split: datasets.Dataset) -> None:
     44     with output_path.open("w") as output_file:
     45         for sample in split:
     46             print(json.dumps({
     47                     "text": input_template(sample),
     48                     "event": output_template(sample),
     49                 }), file=output_file)
     50 
     51 
     52 def main(output_path: pathlib.Path, input_path: pathlib.Path):
     53     output_path.mkdir(parents=True, exist_ok=True)
     54     write_schema(output_path / "event.schema")
     55     for path in input_path.iterdir():
     56         if path.suffix == "":
     57             convert_split(output_path / f"{SPLIT_NAME_MAP[path.name]}.json", datasets.load_from_disk(str(path)))
     58 
     59 
     60 if __name__ == "__main__":
     61     args: argparse.Namespace = parse_command_argument()
     62     main(output_path=args.output, input_path=args.input)