File size: 5,110 Bytes
eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 eb39c91 497d643 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import torch
import transformers
from torch import nn
from torch.nn import CrossEntropyLoss
from typing import Optional, Tuple, Union
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.models.bert.modeling_bert import (
BertPreTrainedModel,
BERT_INPUTS_DOCSTRING,
_TOKENIZER_FOR_DOC,
_CHECKPOINT_FOR_DOC,
BERT_START_DOCSTRING,
_CONFIG_FOR_DOC,
_SEQ_CLASS_EXPECTED_OUTPUT,
_SEQ_CLASS_EXPECTED_LOSS,
BertModel,
)
from transformers.file_utils import (
add_code_sample_docstrings,
add_start_docstrings_to_model_forward,
add_start_docstrings
)
@add_start_docstrings(
"""
Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
output) e.g. for GLUE tasks.
""",
BERT_START_DOCSTRING,
)
class BertForSequenceClassification(BertPreTrainedModel):
def __init__(self, config, **kwargs):
super().__init__(transformers.PretrainedConfig())
#task_labels_map={"binary_classification": 2, "label_classification": 5}
self.tasks = kwargs.get("tasks_map", {})
self.config = config
self.bert = BertModel(config)
classifier_dropout = (
config.classifier_dropout
if config.classifier_dropout is not None
else config.hidden_dropout_prob
)
self.dropout = nn.Dropout(classifier_dropout)
## add task specific output heads
self.classifier1 = nn.Linear(
config.hidden_size, self.tasks[0].num_labels
)
self.classifier2 = nn.Linear(
config.hidden_size, self.tasks[1].num_labels
)
self.init_weights()
@add_start_docstrings_to_model_forward(
BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length")
)
@add_code_sample_docstrings(
processor_class=_TOKENIZER_FOR_DOC,
checkpoint=_CHECKPOINT_FOR_DOC,
output_type=SequenceClassifierOutput,
config_class=_CONFIG_FOR_DOC,
expected_output=_SEQ_CLASS_EXPECTED_OUTPUT,
expected_loss=_SEQ_CLASS_EXPECTED_LOSS,
)
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,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
task_ids=None,
) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
r"""
labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[0, ...,
config.num_labels - 1]`. If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
"""
return_dict = (
return_dict if return_dict is not None else self.config.use_return_dict
)
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,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
pooled_output = outputs[1]
pooled_output = self.dropout(pooled_output)
unique_task_ids_list = torch.unique(task_ids).tolist()
loss_list = []
logits = None
for unique_task_id in unique_task_ids_list:
loss = None
task_id_filter = task_ids == unique_task_id
if unique_task_id == 0:
logits = self.classifier1(pooled_output[task_id_filter])
elif unique_task_id == 1:
logits = self.classifier2(pooled_output[task_id_filter])
if labels is not None:
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.tasks[unique_task_id].num_labels), labels[task_id_filter].view(-1))
loss_list.append(loss)
# logits are only used for eval. and in case of eval the batch is not multi task
# For training only the loss is used
if loss_list:
loss = torch.stack(loss_list).mean()
if not return_dict:
output = (logits,) + outputs[2:]
return ((loss,) + output) if loss is not None else output
return SequenceClassifierOutput(
loss=loss,
logits=logits,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
)
|