File size: 2,500 Bytes
676498f | 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 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from typing import Dict, List
import torch
import torch.nn as nn
from transformers import BertModel, BertTokenizer
class SentimentBERT(nn.Module):
def __init__(self, pretrained_model_name="bert-base-uncased"):
super(SentimentBERT, self).__init__()
self.bert_model = BertModel.from_pretrained(pretrained_model_name)
self.tokenizer = BertTokenizer.from_pretrained(pretrained_model_name)
self.dropout = nn.Dropout(0.3)
self.fc = nn.Linear(self.bert_model.config.hidden_size, 2)
def forward(self, input_ids, attention_mask):
outputs = self.bert_model(input_ids=input_ids, attention_mask=attention_mask)
pooler_output = outputs.pooler_output
output = self.dropout(pooler_output)
output = self.fc(output)
return output
class EndpointHandler:
def __init__(self, path=""):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = SentimentBERT().to(self.device)
# Load the model state dict
if path:
self.model.load_state_dict(torch.load(path, map_location=self.device))
self.model.eval()
self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
def __call__(self, data: Dict) -> Dict:
"""
Args:
data: Dictionary with "inputs" key containing text for sentiment analysis
Returns:
Dictionary with "label" and "score" keys
"""
# Get the input text
inputs = data.pop("inputs", data)
# Ensure the input is a list
if isinstance(inputs, str):
inputs = [inputs]
# Tokenize the input texts
encoded = self.tokenizer(
inputs,
truncation=True,
padding=True,
max_length=512,
return_tensors="pt"
).to(self.device)
# Get model predictions
with torch.no_grad():
outputs = self.model(**encoded)
predictions = torch.softmax(outputs, dim=1)
# Process predictions
results = []
for pred in predictions:
label_id = pred.argmax().item()
score = pred[label_id].item()
label = "Positive" if label_id == 1 else "Negative"
results.append({
"label": label,
"score": round(score, 4)
})
return {"predictions": results}
|