Spaces:
Sleeping
Sleeping
| 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) |