| from dataclasses import dataclass
|
| from typing import Optional, Tuple
|
|
|
| import torch
|
| import torch.nn as nn
|
| from transformers import AutoConfig, AutoModel, PreTrainedModel
|
| from transformers.modeling_outputs import ModelOutput
|
|
|
| from .configuration_suave_multitask import SuaveMultitaskConfig
|
|
|
|
|
| @dataclass
|
| class SuaveMultitaskOutput(ModelOutput):
|
| loss: Optional[torch.FloatTensor] = None
|
| logits_binary: Optional[torch.FloatTensor] = None
|
| logits_multiclass: Optional[torch.FloatTensor] = None
|
| hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
|
|
|
|
|
| class SuaveMultitaskModel(PreTrainedModel):
|
| config_class = SuaveMultitaskConfig
|
| base_model_prefix = "encoder"
|
|
|
| def __init__(self, config: SuaveMultitaskConfig):
|
| super().__init__(config)
|
| base_config = AutoConfig.from_pretrained(config.base_model_name)
|
| self.encoder = AutoModel.from_config(base_config)
|
| hidden_size = self.encoder.config.hidden_size
|
|
|
| self.dropout = nn.Dropout(config.classifier_dropout)
|
| self.classifier_binary = nn.Linear(hidden_size, 2)
|
| self.classifier_multiclass = nn.Linear(hidden_size, config.num_ai_classes)
|
|
|
| self.post_init()
|
|
|
| def forward(
|
| self,
|
| input_ids=None,
|
| attention_mask=None,
|
| labels_binary=None,
|
| labels_multiclass=None,
|
| **kwargs,
|
| ):
|
| outputs = self.encoder(
|
| input_ids=input_ids,
|
| attention_mask=attention_mask,
|
| output_hidden_states=kwargs.get("output_hidden_states", False),
|
| output_attentions=kwargs.get("output_attentions", False),
|
| )
|
|
|
| pooled = outputs.last_hidden_state[:, 0]
|
| pooled = self.dropout(pooled)
|
|
|
| logits_binary = self.classifier_binary(pooled)
|
| logits_multiclass = self.classifier_multiclass(pooled)
|
|
|
| loss = None
|
| if labels_binary is not None and labels_multiclass is not None:
|
| loss_binary = nn.CrossEntropyLoss()(logits_binary, labels_binary)
|
| loss_multiclass = nn.CrossEntropyLoss(ignore_index=-1)(
|
| logits_multiclass, labels_multiclass
|
| )
|
| loss = loss_binary + 0.5 * loss_multiclass
|
|
|
| return SuaveMultitaskOutput(
|
| loss=loss,
|
| logits_binary=logits_binary,
|
| logits_multiclass=logits_multiclass,
|
| hidden_states=outputs.hidden_states,
|
| attentions=outputs.attentions,
|
| )
|
|
|
