ikkbor's picture
Upload folder using huggingface_hub
fba3401 verified
# services/inference.py
import torch
def predict_sentiment(text, tokenizer, model, device):
inputs = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=256,
return_tensors="pt"
)
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
with torch.no_grad():
outputs = model(input_ids, attention_mask)
probs = torch.softmax(outputs, dim=1)
pred = torch.argmax(probs, dim=1).item()
label_map = {0: "Negative", 1: "Positive"}
return label_map[pred], probs[0][pred].item()
def batch_predict(texts, tokenizer, model, device):
inputs = tokenizer(
texts,
padding=True,
truncation=True,
max_length=256,
return_tensors="pt"
)
input_ids = inputs["input_ids"].to(device)
attention_mask = inputs["attention_mask"].to(device)
with torch.no_grad():
outputs = model(input_ids, attention_mask)
probs = torch.softmax(outputs, dim=1)
preds = torch.argmax(probs, dim=1)
label_map = {0: "Negative", 1: "Positive"}
return [
{
"text": text,
"label": label_map[p.item()],
"confidence": probs[i][p].item()
}
for i, (text, p) in enumerate(zip(texts, preds))
]