import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from deep_translator import GoogleTranslator
from langdetect import detect
import torch
import os

MODEL_NAME = "cardiffnlp/twitter-roberta-base-sentiment"
MODEL_DIR = "model"

# Download model if not present
if not os.path.exists(MODEL_DIR) or not os.listdir(MODEL_DIR):
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    tokenizer.save_pretrained(MODEL_DIR)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
    model.save_pretrained(MODEL_DIR)
else:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)

emotion_labels = {
    0: "Negative 😕",
    1: "Neutral 😐",
    2: "Positive 🙂"
}
translator = GoogleTranslator(source='auto', target='en')

def predict_emotion(text):
    detected_language = detect(text)
    if detected_language != 'en':
        translated_text = translator.translate(text)
    else:
        translated_text = text
    inputs = tokenizer(translated_text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=-1).item()
    emotion = emotion_labels.get(predicted_class, "Unknown")
    return emotion

iface = gr.Interface(
    fn=predict_emotion,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here...", label="Input Text"),
    outputs=[
        gr.Textbox(label="Predicted Sentiment")
    ],
    title="Emotion Detection App",
    description="Enter text in any language. The app will detect the language, translate if needed, and predict the emotion."
)

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