dataset.py (5937B)
1 """Module for handling socio-political datasets.""" 2 3 from __future__ import annotations 4 5 import contextlib 6 import datetime 7 import locale 8 import re 9 from typing import TYPE_CHECKING, Any 10 11 import datasets 12 import transformers 13 14 15 if TYPE_CHECKING: 16 import t5.config 17 18 19 class Dataset: 20 """ 21 A socio-political event dataset. 22 23 Config 24 ------ 25 name: str 26 Name of the dataset to train on. It must be a directory in $DATA_PATH. 27 processing: str, optional 28 How to process the data, must be one of "none", "generative". 29 input_template: str, conditional on processing="generative" 30 Template to use for the generative representation of input. 31 output_template: str, conditional on processing="generative" 32 Template to use for the generative representation of output. 33 tokenizer: str 34 Name of the tokenizer used to process text data. 35 36 """ 37 38 def __init__(self: Dataset, config: t5.config.Config) -> None: 39 """Initialize a socio-political event dataset.""" 40 self.config: t5.config.Config = config 41 self.config.declare_option("name", str, "Name of the dataset to train on. It must be a directory in $DATA_PATH.") 42 self.config.declare_option("processing", str, "How to process the data, must be one of \"none\", \"generative\".", default="none") 43 if self.config.processing == "generative": 44 self.config.declare_option("input_template", str, "Template to use for the generative representation of input.") 45 self.config.declare_option("output_template", str, "Template to use for the generative representation of output.") 46 self.config.declare_option("tokenizer", str, "Name of the tokenizer used to process text data.") 47 self.tokenizer: transformers.PreTrainedTokenizer = transformers.AutoTokenizer.from_pretrained(self.config.tokenizer) 48 self._initialize_original_splits() 49 self._initialize_processed_splits() 50 self._initialize_recovery() 51 52 @property 53 def train(self: Dataset) -> datasets.Dataset: 54 """Processed train split.""" 55 return self.datasets["train"] 56 57 @property 58 def validation(self: Dataset) -> datasets.Dataset: 59 """Processed validation split.""" 60 return self.datasets["validation"] 61 62 @property 63 def test(self: Dataset) -> datasets.Dataset: 64 """Processed test split.""" 65 return self.datasets["test"] 66 67 def _initialize_original_splits(self: Dataset) -> None: 68 self.original: dict[str, datasets.Dataset] = {} 69 for directory in (self.config.DATA_PATH / self.config.name).iterdir(): 70 self.original[directory.name] = datasets.load_from_disk(str(directory)) 71 72 def _process_sample_generative(self: Dataset, sample: dict[str, Any]) -> dict[str, Any]: 73 """Process sample as it comes from UCDP into text templates, and tokenize them for training dataset.""" 74 # init huggingface container for sample 75 new_fields: dict[str, Any] = {} 76 77 # parse input & output into plaintext templates from json ucdp sample 78 new_fields["input_text"] = self.config.input_template.format(**sample) 79 new_fields["output_text"] = self.config.output_template.format(**sample, **new_fields) 80 81 # tokenize inputs+outputs & add result to dict 82 new_fields.update({key: value[0] for key, value in self.tokenizer(text=new_fields["input_text"], text_target=new_fields["output_text"], padding="longest", truncation=True, return_tensors="pt").items()}) 83 return new_fields 84 85 def _make_split_batchable(self: Dataset, split: datasets.Dataset) -> datasets.Dataset: 86 return split 87 88 def _process_split(self: Dataset, split: datasets.Dataset) -> datasets.Dataset: 89 if self.config.processing == "none": 90 return split 91 if self.config.processing == "generative": 92 return split.map(self._process_sample_generative) # , load_from_cache_file=False) 93 raise NotImplementedError(f"Unknown value for {self.config.key_path('processing')}: {self.config.processing}") 94 95 def _initialize_processed_splits(self: Dataset) -> None: 96 self.datasets: dict[str, datasets.Dataset] = {} 97 split: str 98 for split in ["train", "validation", "test"]: 99 if split in self.original: 100 self.datasets[split] = self._make_split_batchable(self._process_split(self.original[split])) 101 102 def _initialize_recovery(self: Dataset) -> None: 103 if self.config.processing == "generative": 104 pattern: str = re.sub(r"\s*{[^}]+}\s*", "(.*?)", self.config.output_template) 105 self.template_regex: re.Pattern[str] = re.compile(pattern) 106 self.template_fields: list[str] = re.findall(r"{([^}]+)}", self.config.output_template) 107 108 def recover_data_from_output(self: Dataset, sample: dict[str, Any], output: Any) -> dict[str, Any]: 109 """Given the output of a model, recover the data from which it was generated.""" 110 if self.config.processing == "generative": 111 indices: list[int] = output.tolist() 112 # T5 is using pad as an initial token when generating 113 if indices and indices[0] == self.tokenizer.pad_token_id: 114 indices = indices[1:] 115 with contextlib.suppress(ValueError): 116 # ValueError = no EOS token, this will certainly not match the template regex, try it anyway in case we generated exactly the maximum length. 117 indices = indices[:indices.index(self.tokenizer.eos_token_id)] 118 values: re.Match[str] | None = self.template_regex.fullmatch(self.tokenizer.decode(indices)) 119 if not values: 120 return {} 121 # Convert decoded and matched output template fields to dictionary 122 parsed_values: dict[str, Any] = dict(zip(self.template_fields, (value.strip() for value in values.groups()))) 123 return parsed_values 124 return {}