| import torch | |
| import torch.nn as nn | |
| from transformers import DebertaV2Model, DebertaV2PreTrainedModel | |
| class MultiHeadDebertaForSequenceClassification(DebertaV2PreTrainedModel): | |
| def __init__(self, config, num_heads=5): | |
| super().__init__(config) | |
| self.num_heads = num_heads | |
| self.deberta = DebertaV2Model(config) | |
| self.heads = nn.ModuleList([nn.Linear(config.hidden_size, 4) for _ in range(num_heads)]) | |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) | |
| self.post_init() | |
| def forward(self, input_ids=None, attention_mask=None): | |
| outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask) | |
| sequence_output = outputs[0] | |
| logits_list = [head(self.dropout(sequence_output[:, 0, :])) for head in self.heads] | |
| logits = torch.stack(logits_list, dim=1) | |
| return logits | |