File size: 3,281 Bytes
e77bcc6 |
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 |
from typing import Optional
from transformers import BertModel
from transformers.models.bert.modeling_bert import (
BertPreTrainedModel,
BertOnlyMLMHead,
)
import torch
class BertForPromptFinetuning(BertPreTrainedModel):
def __init__(self, config, use_multi_label_words: bool = False):
super().__init__(config)
self.bert = BertModel(config, add_pooling_layer=False)
self.cls = BertOnlyMLMHead(config)
# Initialize weights and apply final processing
self.init_weights()
self.label_word_ids = None
self.use_multi_label_words = use_multi_label_words
def forward(
self,
input_ids: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
token_type_ids: Optional[torch.Tensor] = None,
mask_pos: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
output_hidden_states: Optional[bool] = False,
output_attentions: Optional[bool] = False,
):
if mask_pos is not None:
mask_pos = mask_pos.squeeze()
elif mask_pos is None:
raise ValueError("`mask_pos` should be assigned!")
# Encode everything
outputs = self.bert(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
output_hidden_states=output_hidden_states,
output_attentions=output_attentions,
)
# Get <mask> token representation
sequence_output = outputs[0]
sequence_mask_output = sequence_output[
torch.arange(sequence_output.size(0)), mask_pos
]
# Logits over vocabulary tokens
# prediction_mask_scores.shape: [batch_size, vocab_size]
prediction_mask_scores = self.cls(sequence_mask_output)
# Return logits for each label
logits = []
if self.use_multi_label_words:
for label_id in self.label_word_ids:
one_label_logits = []
# multiple ids in one label_id
for id in label_id:
one_label_word_logits = prediction_mask_scores[:, id]
one_label_logits.append(one_label_word_logits.unsqueeze(-1))
# one_label_logits: (bs, num_label_words)
one_label_logits = torch.cat(one_label_logits, -1)
# Get the max logits to choose the label word
logits.append(torch.max(one_label_logits, dim=1, keepdim=True)[0])
else:
for label_id in range(len(self.label_word_ids)):
logits.append(
prediction_mask_scores[:, self.label_word_ids[label_id]].unsqueeze(
-1
)
)
# logits.shape: [batch_size, num_classes]
logits = torch.sigmoid(torch.cat(logits, -1))
loss = None
if labels is not None:
loss_fct = torch.nn.BCELoss()
loss = loss_fct(logits, labels.float())
output = (logits, outputs.hidden_states) if output_hidden_states else (logits,)
output = (output + (outputs.attentions)) if output_attentions else output
return ((loss,) + output) if loss is not None else output
|