|
|
| import torch |
| import torch.nn as nn |
| from transformers import AutoModel |
|
|
| class TaskHead(nn.Module): |
| def __init__(self, hidden_size, n_classes, dropout=0.3): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.Linear(hidden_size, 256), |
| nn.LayerNorm(256), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(256, n_classes), |
| ) |
| def forward(self, x): return self.net(x) |
|
|
| class UnifiedMASRIHead(nn.Module): |
| def __init__(self, bert_model_name="T0KII/MASRIBERTv3", ft_dim=300, rnn_hidden=256, num_layers=2, dropout=0.3): |
| super().__init__() |
| self.bert = AutoModel.from_pretrained(bert_model_name) |
| self.bert_hidden = self.bert.config.hidden_size |
| self.bilstm = nn.LSTM(ft_dim, rnn_hidden, num_layers, batch_first=True, bidirectional=True, dropout=dropout) |
| self.bigru = nn.GRU(ft_dim, rnn_hidden, num_layers, batch_first=True, bidirectional=True, dropout=dropout) |
| self.rnn_dropout = nn.Dropout(dropout) |
| combined_dim = self.bert_hidden + (rnn_hidden * 4) |
| self.sarcasm_head = TaskHead(combined_dim, 2, dropout=0.5) |
| self.sentiment_head = TaskHead(combined_dim, 3, dropout=0.3) |
| self.emotion_head = TaskHead(combined_dim, 8, dropout=0.3) |
|
|
| def forward(self, input_ids, attention_mask, ft_embeds, task_name): |
| bert_out = self.bert(input_ids, attention_mask=attention_mask) |
| cls_vec = bert_out.last_hidden_state[:, 0, :] |
| lstm_out, _ = self.bilstm(ft_embeds) |
| gru_out, _ = self.bigru(ft_embeds) |
| rnn_feat = torch.cat([lstm_out[:, -1, :], gru_out[:, -1, :]], dim=1) |
| combined = torch.cat([cls_vec, self.rnn_dropout(rnn_feat)], dim=1) |
| if task_name == 'sarcasm': return self.sarcasm_head(combined) |
| elif task_name == 'sentiment': return self.sentiment_head(combined) |
| elif task_name == 'emotion': return self.emotion_head(combined) |
|
|
