File size: 3,615 Bytes
b5fd657 b9c8e6e b5fd657 8362342 b5fd657 8362342 b5fd657 |
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 |
import torch.nn as nn
from transformers import BertModel, BertPreTrainedModel
from transformers import BertConfig, AutoTokenizer
from .configuration_tiny_greek_news_bert import TinyGreekNewsBertConfig
class TinyGreekNewsBert(BertPreTrainedModel):
config_class = TinyGreekNewsBertConfig
def __init__(self, config):
super().__init__(config)
num_labels_class = config.num_labels_class
num_labels_ner = config.num_labels_ner
self.ner_loss_weight = getattr(config, "ner_loss_weight", 3.0)
self.bert = BertModel(config)
# Classification head
self.class_dropout = nn.Dropout(0.3)
self.class_fc = nn.Linear(config.hidden_size, 768)
self.class_relu = nn.ReLU()
self.classifier = nn.Linear(768, num_labels_class)
# NER head
self.ner_classifier = nn.Linear(config.hidden_size, num_labels_ner)
self.init_weights()
# For normalization
self.initial_cls_loss = None
self.initial_ner_loss = None
def forward(self, input_ids, attention_mask=None, token_type_ids=None,
labels_class=None, labels_ner=None):
outputs = self.bert(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids
)
sequence_output = outputs.last_hidden_state # (batch_size, seq_length, hidden_size)
pooled_output = outputs.pooler_output # (batch_size, hidden_size)
# Classification branch
pooled_output = self.class_dropout(pooled_output)
x = self.class_fc(pooled_output)
x = self.class_relu(x)
logits_class = self.classifier(x)
# NER branch
logits_ner = self.ner_classifier(sequence_output) # (batch_size, seq_length, num_labels_ner)
loss = None
if labels_class is not None and labels_ner is not None:
# Classification loss
loss_fct_class = nn.CrossEntropyLoss()
loss_class = loss_fct_class(logits_class, labels_class)
# NER loss: Cross-entropy with ignore_index=-100, summed then averaged over non-pad tokens
loss_fct_ner = nn.CrossEntropyLoss(ignore_index=-100, reduction='sum')
ner_loss_sum = loss_fct_ner(
logits_ner.view(-1, logits_ner.shape[-1]),
labels_ner.view(-1)
)
mask = (labels_ner != -100).view(-1).float()
loss_ner = ner_loss_sum / (mask.sum() + 1e-9)
# Store initial values
if self.initial_cls_loss is None and self.training:
self.initial_cls_loss = loss_class.item()
if self.initial_ner_loss is None and self.training:
self.initial_ner_loss = loss_ner.item()
# Normalize losses
if (self.initial_cls_loss is not None) and (self.initial_ner_loss is not None):
norm_cls_loss = loss_class / (self.initial_cls_loss + 1e-8)
norm_ner_loss = loss_ner / (self.initial_ner_loss + 1e-8)
else:
norm_cls_loss = loss_class
norm_ner_loss = loss_ner
# Combine losses with weight
loss = norm_cls_loss + self.ner_loss_weight * norm_ner_loss
return (loss, logits_class, logits_ner)
else:
return (logits_class, logits_ner)
TinyGreekNewsBert.register_for_auto_class("AutoModel")
|