|
|
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] |
|
|
pooled = self.dropout(pooled) |
|
|
logits = self.classifier(pooled) |
|
|
return logits |
