|
|
| import torch
|
| import torch.nn as nn
|
| from transformers import PreTrainedModel, BertConfig, BertModel
|
|
|
| def _bert_config_from_base_dict(base_cfg_dict: dict) -> BertConfig:
|
| if base_cfg_dict is None:
|
| raise ValueError("config.base_model_config is required for offline load.")
|
|
|
| base_cfg_dict = dict(base_cfg_dict)
|
| base_cfg_dict["model_type"] = "bert"
|
|
|
| allowed = set(BertConfig().to_dict().keys())
|
| kwargs = {k: v for k, v in base_cfg_dict.items() if k in allowed}
|
|
|
| return BertConfig(**kwargs)
|
|
|
| class KbertMTL(PreTrainedModel):
|
| config_class = BertConfig
|
|
|
| def __init__(self, config):
|
| super().__init__(config)
|
|
|
| base_cfg_dict = getattr(config, "base_model_config", None)
|
| bert_cfg = _bert_config_from_base_dict(base_cfg_dict)
|
| self.bert = BertModel(bert_cfg)
|
|
|
| hidden = self.bert.config.hidden_size
|
| self.head_senti = nn.Linear(hidden, 5)
|
| self.head_act = nn.Linear(hidden, 6)
|
| self.head_emo = nn.Linear(hidden, 7)
|
| self.head_reg = nn.Linear(hidden, 3)
|
| self.has_token_type = getattr(self.bert.embeddings, "token_type_embeddings", None) is not None
|
|
|
| self.post_init()
|
|
|
| def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, **kwargs):
|
| kw = dict(input_ids=input_ids, attention_mask=attention_mask)
|
| if self.has_token_type and token_type_ids is not None:
|
| kw["token_type_ids"] = token_type_ids
|
| out = self.bert(**kw)
|
| h = out.last_hidden_state[:, 0]
|
| return {
|
| "logits_senti": self.head_senti(h),
|
| "logits_act": self.head_act(h),
|
| "logits_emo": self.head_emo(h),
|
| "pred_reg": self.head_reg(h),
|
| "last_hidden_state": out.last_hidden_state
|
| }
|
|
|
