from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import gradio as gr

# Load pre-trained model and tokenizer
model_name = "borisn70/bert-43-multilabel-emotion-detection"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Labels corresponding to different emotions
labels = [
    'admiration', 'amusement', 'anger', 'annoyance', 'approval', 'caring', 'confusion', 
    'curiosity', 'desire', 'disappointment', 'disapproval', 'disgust', 'embarrassment', 
    'excitement', 'fear', 'gratitude', 'grief', 'joy', 'love', 'nervousness', 'optimism', 
    'pride', 'realization', 'relief', 'remorse', 'sadness', 'surprise', 'neutral'
]

# Function to predict emotions based on input text
def predict_emotions(text):
    inputs = tokenizer(text, return_tensors="pt")  # Tokenize the input text
    with torch.no_grad():
        logits = model(**inputs).logits  # Get the model's output logits
    probs = torch.sigmoid(logits)[0]  # Apply sigmoid to get probabilities
    
    # Filter results with probability > 0.5
    results = {label: float(prob) for label, prob in zip(labels, probs) if prob > 0.5}
    return results

# Set up Gradio interface
iface = gr.Interface(fn=predict_emotions, inputs="text", outputs="label")

# Launch the app
iface.launch()