File size: 1,195 Bytes
5225d97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from transformers import RobertaModel, AutoTokenizer
from transformers.modeling_outputs import SequenceClassifierOutput
from huggingface_hub import PyTorchModelHubMixin
from torch.nn import CrossEntropyLoss
import torch.nn as nn
import torch
class SentenceBERTClassifier(nn.Module, PyTorchModelHubMixin):
def __init__(self, model_name="sentence-transformers/all-distilroberta-v1", num_labels=8):
super().__init__()
self.sbert = RobertaModel.from_pretrained(model_name)
self.config = self.sbert.config
self.config.num_labels = num_labels
self.dropout = nn.Dropout(0.05)
self.config.classifier_dropout = 0.05
self.classifier = nn.Linear(self.config.hidden_size, self.config.num_labels)
def forward(self, input_ids, attention_mask):
outputs = self.sbert(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = outputs.pooler_output
dropout_output = self.dropout(pooled_output)
logits = self.classifier(dropout_output)
return SequenceClassifierOutput(
logits=logits,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
) |