import gradio as gr
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer, AutoConfig
from huggingface_hub import hf_hub_download

# Load model and tokenizer
repo_id = "iimran/EmotionDetection"
filename = "model.onnx"

# Download and setup ONNX model
onnx_model_path = hf_hub_download(repo_id=repo_id, filename=filename)
tokenizer = AutoTokenizer.from_pretrained(repo_id)
config = AutoConfig.from_pretrained(repo_id)

# Get label mapping
if hasattr(config, "id2label") and config.id2label and len(config.id2label) > 0:
    id2label = config.id2label
else:
    id2label = {
        0: "anger",
        1: "fear",
        2: "joy",
        3: "love",
        4: "sadness",
        5: "surprise",
        6: "neutral"
    }

# Create ONNX session
session = ort.InferenceSession(onnx_model_path)

def predict_emotion(text):
    """Predict emotion from text"""
    # Tokenize input
    inputs = tokenizer(
        text,
        return_tensors="np",
        truncation=True,
        padding="max_length",
        max_length=256
    )
    
    # Prepare inputs
    ort_inputs = {
        "input_ids": inputs["input_ids"].astype(np.int64),
        "attention_mask": inputs["attention_mask"].astype(np.int64)
    }
    
    # Run inference
    outputs = session.run(None, ort_inputs)
    logits = outputs[0]
    predicted_class_id = int(np.argmax(logits, axis=-1)[0])
    
    # Get label
    predicted_label = id2label.get(str(predicted_class_id), id2label.get(predicted_class_id, str(predicted_class_id)))
    
    # Format output
    emotion_icons = {
        "anger": "😠",
        "fear": "😨",
        "joy": "😄",
        "love": "❤️",
        "sadness": "😢",
        "surprise": "😲",
        "neutral": "😐"
    }
    
    icon = emotion_icons.get(predicted_label.lower(), "❓")
    return f"{icon} {predicted_label}"

# Create Gradio interface
demo = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(label="Enter your text", placeholder="How are you feeling today?"),
    outputs=gr.Label(label="Predicted Emotion"),
    title="Emotion Detection",
    description="Detect emotions in text using iimran/EmotionDetection model",
    examples=[
        ["I'm so happy right now!"],
        ["This situation makes me really angry"],
        ["I feel anxious about the future"],
        ["What a beautiful day to be alive!"],
        ["That news shocked me completely"]
    ],
    theme="soft"
)

# Run the app
if __name__ == "__main__":
    demo.launch()