| 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 |
|
|
