File size: 2,768 Bytes
af4c621 13987ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | # STL
from typing import Optional
import dataclasses
# 3rd Party
import torch
from transformers import BertPreTrainedModel, BertModel
from transformers.models.bert.modeling_bert import BaseModelOutputWithPoolingAndCrossAttentions
from transformers.utils.generic import ModelOutput
# Local
from .configuration_bert_for_stance import BertForStanceConfig
class BertForStance(BertPreTrainedModel):
config_class = BertForStanceConfig
def __init__(self, config: BertForStanceConfig):
super().__init__(config)
self.config = config
self.num_labels = config.num_labels
self.bert = BertModel(config)
hidden_size = config.hidden_size
classifier_hidden_units = config.classifier_hidden_units or config.hidden_size
self.classifier = torch.nn.Sequential(
torch.nn.Dropout(config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob),
torch.nn.Linear(hidden_size, classifier_hidden_units, bias=True),
torch.nn.ReLU(),
torch.nn.Linear(classifier_hidden_units, self.num_labels, bias=True)
)
self.loss_fct = torch.nn.CrossEntropyLoss()
self.post_init()
@dataclasses.dataclass
class Output(ModelOutput):
loss: Optional[torch.FloatTensor] = None
logits: Optional[torch.FloatTensor] = None
seq_encoding: Optional[torch.FloatTensor] = None
def forward(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
token_type_ids: Optional[torch.Tensor] = None,
position_ids: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
return_dict: Optional[bool] = None,
):
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
outputs: BaseModelOutputWithPoolingAndCrossAttentions = self.bert(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
output_attentions=False,
output_hidden_states=True,
return_dict=True,
)
feature_vec = outputs.last_hidden_state[:, 0]
logits = self.classifier(feature_vec)
loss = None
if labels is not None:
loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
return BertForStance.Output(loss=loss, logits=logits, seq_encoding=feature_vec)
BertForStance.register_for_auto_class("AutoModel") |