Update modeling.py

modeling.py

@@ -1,64 +1,22 @@
 import torch
-import gradio as gr
-from transformers import AutoTokenizer
-
-LABELS = ["anger", "fear", "joy", "sadness", "surprise"]   # your 5 labels
-
-tokenizer = AutoTokenizer.from_pretrained("./")
-
-model = BERTMultiLabel(
-    model_name="microsoft/deberta-v3-base",
-    num_labels=len(LABELS)
-)
-
-state = torch.load("pytorch_model.bin", map_location="cpu")
-model.load_state_dict(state)
-model.eval()
-
-def predict(text):
-    if not text.strip():
-        return {"error": "Please enter text"}
-
-    enc = tokenizer(
-        text,
-        truncation=True,
-        padding="max_length",
-        max_length=128,
-        return_tensors="pt"
-    )
-
-    with torch.no_grad():
-        logits = model(enc["input_ids"], enc["attention_mask"])
-        probs = torch.sigmoid(logits)[0].tolist()
-
-    # probability dictionary
-    scores = {label: round(p, 4) for label, p in zip(LABELS, probs)}
-
-    # main mood (highest probability)
-    main_emotion = LABELS[int(torch.tensor(probs).argmax())]
-
-    return {
-        "Predicted Mood": main_emotion,
-        "Scores": scores
-    }
-
-with gr.Blocks() as app:
-    gr.Markdown("# 🧠 Mood Detection using DeBERTa-v3")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("""
-            ## 📌 Model Info
-            - **Model:** DeBERTa-v3 Base  
-            - **Task:** Multi-label Sentiment / Emotion  
-            - **Framework:** PyTorch + Transformers  
-            - **Labels:** anger, fear, joy, sadness, surprise  
-            """)
-        
-        with gr.Column(scale=2):
-            user_input = gr.Textbox(label="Enter text here")
-            btn = gr.Button("Analyze Mood")
-            output = gr.JSON(label="Output")
-            btn.click(fn=predict, inputs=user_input, outputs=output)
-
-app.launch()
+import torch.nn as nn
+from transformers import AutoModel
+
+class BERTMultiLabel(nn.Module):
+    def __init__(self, model_name="microsoft/deberta-v3-base", num_labels=5):
+        super().__init__()
+        self.bert = AutoModel.from_pretrained(model_name)
+        hidden = self.bert.config.hidden_size
+        self.dropout = nn.Dropout(0.2)
+        self.classifier = nn.Linear(hidden, num_labels)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.bert(
+            input_ids=input_ids,
+            attention_mask=attention_mask
+        )
+
+        cls = outputs.last_hidden_state[:, 0]  # CLS token
+        cls = self.dropout(cls)
+        logits = self.classifier(cls)
+        return logits