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)