inference.py

from transformers import BertForSequenceClassification, BertTokenizer
import torch
import json
import requests

def predict(text):
    repo_id = "logasanjeev/goemotions-bert"
    model = BertForSequenceClassification.from_pretrained(repo_id)
    tokenizer = BertTokenizer.from_pretrained(repo_id)
    thresholds_url = f"https://huggingface.co/{repo_id}/raw/main/thresholds.json"
    thresholds_data = json.loads(requests.get(thresholds_url).text)
    emotion_labels = thresholds_data["emotion_labels"]
    thresholds = thresholds_data["thresholds"]
    encodings = tokenizer(text, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
    with torch.no_grad():
        logits = torch.sigmoid(model(**encodings).logits).numpy()[0]
    predictions = [{"label": emotion_labels[i], "score": float(logit)} for i, (logit, thresh) in enumerate(zip(logits, thresholds)) if logit >= thresh]
    return sorted(predictions, key=lambda x: x["score"], reverse=True)