|
from typing import Optional |
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from torch import nn |
|
from transformers import AutoModel |
|
from transformers import PreTrainedModel |
|
|
|
from .configuration_leaf import LeafConfig |
|
from .mappings import idx_to_ef, idx_to_classname |
|
|
|
|
|
class LeafModel(PreTrainedModel): |
|
""" |
|
LEAF model for text classification. |
|
""" |
|
config_class = LeafConfig |
|
|
|
def __init__(self, config: LeafConfig): |
|
super().__init__(config) |
|
self._base_model = AutoModel.from_pretrained(config.model_name) |
|
self._device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
hidden_dim = self._base_model.config.hidden_size |
|
self.head = ClassificationHead(hidden_dim=hidden_dim, num_classes=2097, |
|
idx_to_ef=idx_to_ef, idx_to_classname=idx_to_classname, |
|
device=self._device) |
|
|
|
def forward(self, input_ids, attention_mask, **kwargs) -> dict: |
|
if "classes" not in kwargs: |
|
kwargs["classes"] = None |
|
outputs = self._base_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state |
|
attention_mask = attention_mask.unsqueeze(-1) |
|
masked_outputs = outputs * attention_mask.type_as(outputs) |
|
nom = masked_outputs.sum(dim=1) |
|
denom = attention_mask.sum(dim=1) |
|
denom = denom.masked_fill(denom == 0, 1) |
|
return self.head(nom / denom, **kwargs) |
|
|
|
|
|
class ClassificationHead(nn.Module): |
|
""" |
|
Model head to predict a categorical target variable. |
|
""" |
|
|
|
def __init__(self, hidden_dim: int, num_classes: int, idx_to_ef: dict, idx_to_classname: Optional[dict], |
|
device: str): |
|
super().__init__() |
|
self.linear = nn.Linear(in_features=hidden_dim, out_features=num_classes) |
|
self.loss = nn.CrossEntropyLoss() |
|
|
|
|
|
self.idx_to_ef = torch.Tensor([idx_to_ef[k] for k in sorted(idx_to_ef.keys())]).to(device) |
|
self.idx_to_ef.requires_grad = False |
|
self.idx_to_classname = idx_to_classname |
|
|
|
def __call__(self, activations: torch.Tensor, classes: Optional[torch.Tensor], **kwargs) -> dict: |
|
return_dict = {} |
|
logits = self.linear(activations) |
|
return_dict["logits"] = logits |
|
if classes: |
|
loss = self.loss(logits, classes) |
|
return_dict["loss"] = loss |
|
_, predicted_classes = torch.max(F.softmax(logits, dim=1), dim=1) |
|
return_dict["class_idx"] = predicted_classes |
|
return_dict["ef_score"] = self.idx_to_ef[predicted_classes] |
|
if self.idx_to_classname: |
|
return_dict["class"] = [self.idx_to_classname[str(c)] for c in |
|
predicted_classes.cpu().numpy()] |
|
return return_dict |
|
|
