File size: 922 Bytes
9434e7d |
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 |
from transformers import AutoModelForSequenceClassification
import torch
class FineTunedDistilBertWithStringLabels(AutoModelForSequenceClassification):
def __init__(self, config):
super().__init__(config)
self.label_dict = {
0: "Clarification",
1: "Factual",
2: "Operational",
3: "Summarization"
}
def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None):
# Perform the usual forward pass
outputs = super().forward(input_ids, attention_mask, token_type_ids, labels)
# Get logits (raw model output)
logits = outputs.logits
predicted_class_index = torch.argmax(logits, dim=-1).item() # .item() -> extract regular python number
# Map index -> string label using dictionary
outputs.label = self.label_dict[predicted_class_index]
return outputs
|