Views
No views yet
1from transformers import AutoModelForSequenceClassification, AutoTokenizer
2import torch
3import numpy as np
4
5def predict_emotions(text, model_name, threshold=0.35):
6 # Load model and tokenizer
7 model = AutoModelForSequenceClassification.from_pretrained(model_name)
8 tokenizer = AutoTokenizer.from_pretrained(model_name)
9
10 # Tokenize and predict
11 inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=250)
12 with torch.no_grad():
13 outputs = model(**inputs)
14 logits = outputs.logits
15 probabilities = torch.sigmoid(logits).numpy()[0]
16
17 # Map probabilities to emotions
18 emotions = {emotion: float(prob) for emotion, prob in zip(model.config.id2label.values(), probabilities)}
19
20 # Get emotions above threshold and sort by probability
21 predicted_emotions = [(emotion, prob) for emotion, prob in emotions.items() if prob >= threshold]
22 predicted_emotions.sort(key=lambda x: x[1], reverse=True)
23
24 return {
25 "text": text,
26 "predicted_emotions": predicted_emotions,
27 "all_probabilities": dict(sorted(emotions.items(), key=lambda x: x[1], reverse=True)),
28 "threshold_used": threshold
29 }
30
31# Example usage
32result = predict_emotions(
33 "I'm feeling really excited and happy about this news!",
34 "model-name",
35 threshold=0.35 # Customize threshold here
36)
37
38# Print results
39print(f"Text: {result['text']}")
40print("\nDetected emotions (sorted by probability):")
41for emotion, prob in result['predicted_emotions']:
42 print(f" - {emotion.upper()} ({prob:.4f})")
43
44print("\nAll emotion probabilities (sorted):")
45for emotion, prob in result['all_probabilities'].items():
46 print(f" {'*' if prob >= result['threshold_used'] else ' '} {emotion}: {prob:.4f}")