File size: 5,703 Bytes
e61b0cc | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig, AutoModel, AutoConfig
class MultiHeadConfig(PretrainedConfig):
model_type = "multihead_text_classifier"
def __init__(self, encoder_name="cointegrated/LaBSE-en-ru",
num_labels_l2=5, num_labels_l3=11, num_labels_l4=41,
dropout=0.1, hidden_size=768,
id2label_l2=None, label2id_l2=None,
id2label_l3=None, label2id_l3=None,
id2label_l4=None, label2id_l4=None,
l3_to_l2=None, l4_to_l3=None, merged_labels=None,
loss_weight_l2=0.25, loss_weight_l3=0.5, loss_weight_l4=1.0,
label_smoothing=0.0, **kwargs):
super().__init__(**kwargs)
self.encoder_name = encoder_name
self.num_labels_l2 = num_labels_l2
self.num_labels_l3 = num_labels_l3
self.num_labels_l4 = num_labels_l4
self.dropout = dropout
self.hidden_size = hidden_size
self.id2label_l2 = id2label_l2 or {}
self.label2id_l2 = label2id_l2 or {}
self.id2label_l3 = id2label_l3 or {}
self.label2id_l3 = label2id_l3 or {}
self.id2label_l4 = id2label_l4 or {}
self.label2id_l4 = label2id_l4 or {}
self.l3_to_l2 = l3_to_l2 or [] # index = l3_id -> l2_id
self.l4_to_l3 = l4_to_l3 or [] # index = l4_id -> l3_id; -1 у мёрж-классов TC1/TC2
self.merged_labels = merged_labels or [] # метки, требующие слоя правил (TC1/TC2)
self.loss_weight_l2 = loss_weight_l2
self.loss_weight_l3 = loss_weight_l3
self.loss_weight_l4 = loss_weight_l4
self.label_smoothing = label_smoothing
class MultiHeadClassifier(PreTrainedModel):
"""LaBSE -> [CLS] -> dropout -> три линейные головы (L2 / L3 / L4)."""
config_class = MultiHeadConfig
def __init__(self, config):
super().__init__(config)
enc_config = AutoConfig.from_pretrained(config.encoder_name)
self.encoder = AutoModel.from_config(enc_config) # только архитектура (веса грузятся отдельно)
h = enc_config.hidden_size
self.dropout = nn.Dropout(config.dropout)
self.head_l2 = nn.Linear(h, config.num_labels_l2)
self.head_l3 = nn.Linear(h, config.num_labels_l3)
self.head_l4 = nn.Linear(h, config.num_labels_l4)
# веса классов (не сохраняются в чекпойнт): по умолчанию единицы
self.register_buffer("cw_l2", torch.ones(config.num_labels_l2), persistent=False)
self.register_buffer("cw_l3", torch.ones(config.num_labels_l3), persistent=False)
self.register_buffer("cw_l4", torch.ones(config.num_labels_l4), persistent=False)
self.post_init()
@classmethod
def from_encoder(cls, config):
"""Инициализация для обучения: подгружаем претренированные веса энкодера."""
model = cls(config)
model.encoder = AutoModel.from_pretrained(config.encoder_name)
return model
def set_class_weights(self, w_l2, w_l3, w_l4):
self.register_buffer("cw_l2", torch.as_tensor(w_l2, dtype=torch.float), persistent=False)
self.register_buffer("cw_l3", torch.as_tensor(w_l3, dtype=torch.float), persistent=False)
self.register_buffer("cw_l4", torch.as_tensor(w_l4, dtype=torch.float), persistent=False)
def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
labels_l2=None, labels_l3=None, labels_l4=None, **kwargs):
out = self.encoder(input_ids=input_ids, attention_mask=attention_mask,
token_type_ids=token_type_ids)
cls = out.last_hidden_state[:, 0] # [CLS]
cls = self.dropout(cls)
logits_l2 = self.head_l2(cls)
logits_l3 = self.head_l3(cls)
logits_l4 = self.head_l4(cls)
loss = None
if labels_l2 is not None and labels_l3 is not None and labels_l4 is not None:
ls = getattr(self.config, "label_smoothing", 0.0)
def _ce(logits, labels, cw):
# ignore_index=-100: у мёрж-классов (TC1/TC2) нет однозначного L2/L3, и лосс этих
# голов на таких строках выключается маркером -100 в labels (см. Шаг 3).
w = cw.to(logits.device).to(logits.dtype)
# если ВЕСЬ батч замаскирован, cross_entropy вернул бы NaN (деление на нулевую
# сумму весов) и отравил бы обучение -> отдаём ровно нулевой вклад, сохраняя граф.
if int((labels != -100).sum()) == 0:
return logits.sum() * 0.0
return F.cross_entropy(logits, labels, weight=w, label_smoothing=ls,
ignore_index=-100)
l2 = _ce(logits_l2, labels_l2, self.cw_l2)
l3 = _ce(logits_l3, labels_l3, self.cw_l3)
l4 = _ce(logits_l4, labels_l4, self.cw_l4)
loss = (self.config.loss_weight_l2 * l2
+ self.config.loss_weight_l3 * l3
+ self.config.loss_weight_l4 * l4)
return {"loss": loss, "logits_l2": logits_l2, "logits_l3": logits_l3, "logits_l4": logits_l4}
MultiHeadConfig.register_for_auto_class()
MultiHeadClassifier.register_for_auto_class("AutoModel")
|