Update app.py

app.py

@@ -1,9 +1,19 @@
 import gradio as gr
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+# Load the model and tokenizer once during initialization
+model_name = "AnkitAI/deberta-xlarge-base-emotions-classifier"
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Define the function to use the model for predictions
 def classify_emotion(text):
-    model = gr.Interface.load("models/AnkitAI/deberta-xlarge-base-emotions-classifier")
-    return model(text)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs)
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    labels = ["joy", "anger", "sadness", "fear", "surprise", "love"]  # Adjust based on the actual labels used by the model
+    return {labels[i]: float(probs[0][i]) for i in range(len(labels))}
 
 # Validate the input
 def validate_input(text):
@@ -19,7 +29,6 @@ interface = gr.Interface(
     title="Emotion Classifier",
     description="Enter some text and let the model predict the emotion.",
     examples=["I am feeling great today!", "I am so sad and depressed.", "I am excited about the new project."],
-    theme="huggingface"
 )
 
 # Add some custom CSS to improve the look and feel