PraneshJs's picture
Update gri.py
cf53d89 verified
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)