зеркало из https://github.com/microsoft/MWSS.git
93 строки
4.0 KiB
Python
93 строки
4.0 KiB
Python
from transformers import BertPreTrainedModel, BertModel
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class BertForSequenceClassification(BertPreTrainedModel):
|
|
r"""
|
|
**labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:
|
|
Labels for computing the sequence classification/regression loss.
|
|
Indices should be in ``[0, ..., config.num_labels - 1]``.
|
|
If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
|
|
If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
|
|
|
|
Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
|
|
**loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
|
|
Classification (or regression if config.num_labels==1) loss.
|
|
**logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
|
|
Classification (or regression if config.num_labels==1) scores (before SoftMax).
|
|
**hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
|
|
list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
|
|
of shape ``(batch_size, sequence_length, hidden_size)``:
|
|
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
|
|
**attentions**: (`optional`, returned when ``config.output_attentions=True``)
|
|
list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
|
|
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
|
|
|
|
Examples::
|
|
|
|
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
|
|
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
|
|
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0) # Batch size 1
|
|
labels = torch.tensor([1]).unsqueeze(0) # Batch size 1
|
|
outputs = model(input_ids, labels=labels)
|
|
loss, logits = outputs[:2]
|
|
|
|
"""
|
|
|
|
def __init__(self, config):
|
|
super(BertForSequenceClassification, self).__init__(config)
|
|
self.num_labels = config.num_labels
|
|
|
|
self.bert = BertModel(config)
|
|
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
|
self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
|
|
|
|
self.init_weights()
|
|
|
|
def expand_class_head(self, c_count):
|
|
self.c_count = c_count
|
|
if c_count > 1:
|
|
for i in range(1, c_count + 1):
|
|
setattr(self, "classifier_{}".format(i), nn.Linear(self.config.hidden_size, self.num_labels))
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
attention_mask=None,
|
|
token_type_ids=None,
|
|
position_ids=None,
|
|
head_mask=None,
|
|
inputs_embeds=None,
|
|
labels=None,
|
|
reduction="mean", is_gold=True,
|
|
|
|
):
|
|
|
|
outputs = 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,
|
|
)
|
|
|
|
pooled_output = outputs[1]
|
|
outputs = (pooled_output.detach(), )
|
|
pooled_output = self.dropout(pooled_output)
|
|
|
|
if self.c_count == 1 or is_gold:
|
|
logits = self.classifier(pooled_output) # (N, C)
|
|
else:
|
|
logits = []
|
|
for i in range(1, self.c_count+1):
|
|
logits.append(self.__dict__['_modules']["classifier_{}".format(i)](pooled_output))
|
|
logits = torch.cat(logits, dim=1)
|
|
|
|
outputs = (logits, ) + outputs
|
|
if labels is not None:
|
|
loss_fct = nn.CrossEntropyLoss(reduction=reduction)
|
|
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
|
|
outputs = (loss,) + outputs
|
|
|
|
return outputs # (loss), logits, (hidden_states), (attentions) |