| | |
| | from __future__ import annotations |
| |
|
| | from dataclasses import dataclass |
| | from typing import Optional |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from transformers import AutoConfig, AutoModel, PreTrainedModel |
| | from transformers.utils import ModelOutput |
| |
|
| |
|
| | def hidden_size_from_config(config) -> int: |
| | return int(getattr(config, "hidden_size", getattr(config, "dim"))) |
| |
|
| |
|
| | @dataclass |
| | class MultilabelSpanOutput(ModelOutput): |
| | loss: Optional[torch.Tensor] = None |
| | start_logits: Optional[torch.Tensor] = None |
| | end_logits: Optional[torch.Tensor] = None |
| |
|
| |
|
| | class IrishCoreSpanHeadModel(PreTrainedModel): |
| | config_class = AutoConfig |
| | base_model_prefix = "encoder" |
| |
|
| | def __init__(self, config): |
| | super().__init__(config) |
| | num_span_labels = int(getattr(config, "num_span_labels")) |
| | self.encoder = AutoModel.from_config(config) |
| | hidden_size = hidden_size_from_config(config) |
| | dropout = float(getattr(config, "seq_classif_dropout", getattr(config, "dropout", 0.1))) |
| | self.dropout = nn.Dropout(dropout) |
| | self.start_classifier = nn.Linear(hidden_size, num_span_labels) |
| | self.end_classifier = nn.Linear(hidden_size, num_span_labels) |
| | pos_weight = float(getattr(config, "span_positive_weight", 6.0)) |
| | self.register_buffer("loss_pos_weight", torch.full((num_span_labels,), pos_weight), persistent=False) |
| | self.post_init() |
| |
|
| | def forward( |
| | self, |
| | input_ids=None, |
| | attention_mask=None, |
| | token_type_ids=None, |
| | start_positions=None, |
| | end_positions=None, |
| | token_mask=None, |
| | **kwargs, |
| | ) -> MultilabelSpanOutput: |
| | encoder_kwargs = { |
| | "input_ids": input_ids, |
| | "attention_mask": attention_mask, |
| | **kwargs, |
| | } |
| | if token_type_ids is not None and getattr(self.config, "model_type", "") not in {"distilbert", "roberta"}: |
| | encoder_kwargs["token_type_ids"] = token_type_ids |
| | outputs = self.encoder(**encoder_kwargs) |
| | hidden = self.dropout(outputs.last_hidden_state) |
| | start_logits = self.start_classifier(hidden) |
| | end_logits = self.end_classifier(hidden) |
| |
|
| | loss = None |
| | if start_positions is not None and end_positions is not None: |
| | if token_mask is None: |
| | token_mask = attention_mask |
| | mask = token_mask.float().unsqueeze(-1) |
| | pos_weight = self.loss_pos_weight.to(start_logits.device) |
| | bce = nn.BCEWithLogitsLoss(reduction="none", pos_weight=pos_weight) |
| | start_loss = bce(start_logits, start_positions.float()) * mask |
| | end_loss = bce(end_logits, end_positions.float()) * mask |
| | denom = mask.sum().clamp_min(1.0) * start_logits.shape[-1] |
| | loss = (start_loss.sum() + end_loss.sum()) / (2.0 * denom) |
| |
|
| | return MultilabelSpanOutput(loss=loss, start_logits=start_logits, end_logits=end_logits) |
| |
|
