File size: 665 Bytes
8beae03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel
class MultiLabelDeberta(nn.Module):
def __init__(self, num_labels):
super().__init__()
self.backbone = AutoModel.from_pretrained('microsoft/deberta-v3-base')
self.dropout = nn.Dropout(0.3)
self.classifier = nn.Linear(self.backbone.config.hidden_size, num_labels)
def forward(self, input_ids, attention_mask):
outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
pooled = outputs.last_hidden_state[:, 0] # [CLS]
pooled = self.dropout(pooled)
logits = self.classifier(pooled)
return logits |