import gradio as gr
import torch
from peft import PeftModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import json

# Load model and tokenizer
base_model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=6)
model = PeftModel.from_pretrained(base_model, "katsuchi/bert-dair-ai-emotion")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

def predict_emotion(text):
    # Tokenize input
    tokens = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    
    # Get model prediction
    with torch.no_grad():
        outputs = model(tokens['input_ids'])
        probs = torch.softmax(outputs.logits, dim=-1)
    
    # Convert probabilities to percentages
    percentages = (probs * 100).squeeze().tolist()
    
    # Create emotion-percentage mapping
    emotions = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
    emotion_probs = {
        emotion: f"{percentage:.1f}%" 
        for emotion, percentage in zip(emotions, percentages)
    }
    
    # Sort by probability in descending order
    sorted_emotions = dict(
        sorted(emotion_probs.items(), 
               key=lambda x: float(x[1].rstrip('%')), 
               reverse=True)
    )
    
    # Format output
    return json.dumps(sorted_emotions, indent=2)

# Create Gradio interface
iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(
        lines=3, 
        placeholder="Enter text here..."
    ),
    outputs=gr.JSON(),
    title="Emotion Classifier",
    description="Predict emotions in text with confidence percentages",
    examples=[
        ["I am so happy to see you!"],
        ["I'm really disappointed with the results."],
        ["That's absolutely terrifying!"],
        ["I love spending time with my family."]
    ]
)

if __name__ == "__main__":
    iface.launch()