Update app.py

app.py

@@ -1,46 +1,71 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
-
-# Load models
-model_paths = {
-    "BERT": "models/bert_model",
-    "XLNet": "models/xlnet_model",
-    "GPT-2": "models/gpt2_model"
-}
-
-models = {}
-tokenizers = {}
-
-for name, path in model_paths.items():
-    tokenizers[name] = AutoTokenizer.from_pretrained(path)
-    models[name] = AutoModelForSequenceClassification.from_pretrained(path)
-    models[name].eval()
-
-# Emotion labels (adjust based on your dataset!)
-labels = ["anger", "joy", "sadness", "fear", "love", "surprise"]
-
-def classify(text, model_choice):
-    tokenizer = tokenizers[model_choice]
-    model = models[model_choice]
-
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-        top_prob, top_idx = torch.max(probs, dim=1)
-
-    return f"Predicted: {labels[top_idx.item()]} ({top_prob.item():.2f})"
-
-iface = gr.Interface(
-    fn=classify,
-    inputs=[
-        gr.Textbox(lines=3, placeholder="Enter text here...", label="Text"),
-        gr.Radio(choices=["BERT", "XLNet", "GPT-2"], label="Choose Model")
-    ],
-    outputs="text",
-    title="Emotion Classifier (BERT / XLNet / GPT-2)"
-)
-
-if __name__ == "__main__":
-    iface.launch()
+import gradio as gr
+import torch
+from transformers import (
+    BertTokenizer,
+    XLNetTokenizer,
+    GPT2Tokenizer,
+    AutoModelForSequenceClassification
+)
+
+# Model repositories on Hugging Face Hub
+model_repos = {
+    "BERT": "sk23aib/emotion-bert",
+    "XLNet": "sk23aib/emotion-xlnet",
+    "GPT-2": "sk23aib/emotion-gpt2"
+}
+
+# Emotion labels (must match model training order)
+emotion_labels = [
+    "anger", "boredom", "empty", "enthusiasm", "fun", "happiness", "hate",
+    "love", "neutral", "relief", "sadness", "surprise", "worry"
+]
+
+# Load models and tokenizers
+loaded_models = {}
+
+# BERT
+bert_tokenizer = BertTokenizer.from_pretrained(model_repos["BERT"])
+bert_model = AutoModelForSequenceClassification.from_pretrained(model_repos["BERT"])
+bert_model.eval()
+loaded_models["BERT"] = {"tokenizer": bert_tokenizer, "model": bert_model}
+
+# XLNet
+xlnet_tokenizer = XLNetTokenizer.from_pretrained(model_repos["XLNet"])
+xlnet_model = AutoModelForSequenceClassification.from_pretrained(model_repos["XLNet"])
+xlnet_model.eval()
+loaded_models["XLNet"] = {"tokenizer": xlnet_tokenizer, "model": xlnet_model}
+
+# GPT-2
+gpt2_tokenizer = GPT2Tokenizer.from_pretrained(model_repos["GPT-2"], padding_side="left")
+gpt2_tokenizer.pad_token = gpt2_tokenizer.eos_token  # Required for GPT-2
+gpt2_model = AutoModelForSequenceClassification.from_pretrained(model_repos["GPT-2"])
+gpt2_model.config.pad_token_id = gpt2_tokenizer.pad_token_id
+gpt2_model.eval()
+loaded_models["GPT-2"] = {"tokenizer": gpt2_tokenizer, "model": gpt2_model}
+
+# Inference function – return top emotion + probability
+def predict_emotions(text):
+    output_lines = []
+    with torch.no_grad():
+        for model_name, components in loaded_models.items():
+            tokenizer = components["tokenizer"]
+            model = components["model"]
+            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+            logits = model(**inputs).logits
+            probs = torch.nn.functional.softmax(logits, dim=1)[0]
+            top_idx = torch.argmax(probs).item()
+            top_emotion = emotion_labels[top_idx]
+            top_confidence = round(float(probs[top_idx]), 4)
+            output_lines.append(f"{model_name}: {top_emotion} ({top_confidence})")
+    return "\n".join(output_lines)
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=predict_emotions,
+    inputs=gr.Textbox(lines=3, placeholder="Type a sentence to analyze..."),
+    outputs=gr.Textbox(label="Top Emotion by Model"),
+    title="Multi-Model Emotion Classifier",
+    description="See which emotion is predicted by BERT, XLNet, and GPT-2, along with their confidence."
+)
+
+interface.launch()