model.py (3282B)
1 """Module for handling socio-political models.""" 2 3 from __future__ import annotations 4 5 from typing import TYPE_CHECKING, Any 6 7 import torch 8 import transformers 9 from transformers import EvalPrediction 10 11 import t5.config 12 import t5.dataset 13 14 15 class ModelConfig(transformers.PretrainedConfig): 16 """Binding class between huggingface config and our own.""" 17 18 def __init__(self: ModelConfig, config: t5.config.Config | dict[str, dict[str, str]] | None = None, **kwargs: Any) -> None: 19 """Initialize a ModelConfig either from a config or from a loaded json object.""" 20 if isinstance(config, t5.config.Config): 21 self.config = config 22 elif isinstance(config, dict): 23 self.config = t5.config.Config(root=["model"], lazy_values = config["config"]) 24 else: 25 self.config = t5.config.Config(root=["model"]) 26 super().__init__(**kwargs) 27 28 def to_dict(self: ModelConfig) -> dict[str, Any]: 29 """Convert the configuration to a json object using strings for all values.""" 30 return {"config": { 31 key: str(value) 32 for key, value in self.config.items() 33 }} 34 35 36 class Model(transformers.PreTrainedModel): 37 """ 38 A socio-political event extraction model. 39 40 Config 41 ------ 42 transformer: str 43 Name of the huggingface model to use. 44 45 generation_max_length: int 46 Maximum number of tokens to generate during evaluetion. 47 48 """ 49 50 config_class = ModelConfig 51 52 def __init__(self: Model, config: t5.config.Config, _: t5.dataset.Dataset) -> None: 53 """Initialize a socio-political event extraction model.""" 54 self.config = ModelConfig(config) 55 super().__init__(self.config) 56 config.declare_option("transformer", str, "Name of the huggingface model to use.") 57 self.seq2seq = transformers.AutoModelForSeq2SeqLM.from_pretrained(config.transformer) 58 config.declare_option("generation_max_length", int, "Maximum number of tokens to generate during evaluetion.") 59 self.seq2seq.generation_config.max_length = config.generation_max_length 60 self._keys_to_ignore_on_save = None 61 62 def forward(self: Model, input_ids: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor, **kwargs: Any) -> transformers.modeling_outputs.Seq2SeqLMOutput: 63 """Return the loss for a given input and gold output.""" 64 del kwargs 65 output: transformers.modeling_outputs.Seq2SeqLMOutput = self.seq2seq(input_ids=input_ids, attention_mask=attention_mask, labels=labels) 66 return output 67 68 def generate(self: Model, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs: Any) -> torch.Tensor: 69 """Extract the information for a given input.""" 70 del kwargs 71 output: torch.Tensor = self.seq2seq.generate(input_ids=input_ids, attention_mask=attention_mask) 72 return output 73 74 @staticmethod 75 def detect_prediction_generation(output: EvalPrediction) -> Any: 76 """Detect if the model output was predicted with generate() or not.""" 77 predictions = output.predictions 78 if isinstance(predictions, torch.Tensor): 79 return predictions 80 if isinstance(predictions, tuple): 81 return torch.argmax(predictions[0], dim=-1)