| import torch |
| from torch import nn |
| from transformers import DistilBertModel, DistilBertTokenizer |
|
|
| class MultilabelClassifier(nn.Module): |
| """Base model for multilabel classification supporting different backbones""" |
| |
| def __init__(self, model_name, num_labels, dropout=0.1): |
| super(MultilabelClassifier, self).__init__() |
| |
| self.backbone = DistilBertModel.from_pretrained(model_name) |
| self.dropout = nn.Dropout(dropout) |
| self.classifier = nn.Linear(768, num_labels) |
| self.sigmoid = nn.Sigmoid() |
| |
| def forward(self, input_ids, attention_mask): |
| outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask) |
| pooled_output = outputs.last_hidden_state[:, 0] |
| x = self.dropout(pooled_output) |
| logits = self.classifier(x) |
| return self.sigmoid(logits) |
|
|
