Update app.py

app.py

@@ -1,91 +1,100 @@
-import json
-from pathlib import Path
-from typing import Dict, Tuple, List
-
 import gradio as gr
 import torch
 import torch.nn as nn
+import json
 from transformers import AutoTokenizer, AutoConfig, AutoModel
 
+# ----------------------------
+# Configuration / Meta
+# ----------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-META_PATH = "meta.json"
-WEIGHTS_PATH = "best_model.pth"
+MAX_LENGTH = 256
+WEIGHTS_PATH = "best_model.pth"  # your model file
 TOKENIZER_PATH = "tokenizer"
+META_PATH = "meta.json"
 
+# Load meta.json
+with open(META_PATH, "r") as f:
+    meta = json.load(f)
+TARGET_COLS = meta.get("target_cols", ["anger","fear","joy","sadness","surprise"])
+MAX_LENGTH = meta.get("max_length", MAX_LENGTH)
+BACKBONE = meta.get("backbone", "microsoft/deberta-v3-base")
 
-def load_meta(path: str) -> Dict:
-    p = Path(path)
-    if not p.exists():
-        return {"target_cols": ["anger", "fear", "joy", "sadness", "surprise"], "max_length": 256, "backbone": "microsoft/deberta-v3-base"}
-    with p.open("r", encoding="utf8") as f:
-        return json.load(f)
-
-
-meta = load_meta(META_PATH)
-TARGET_COLS: List[str] = meta.get("target_cols", ["anger", "fear", "joy", "sadness", "surprise"])
-MAX_LENGTH: int = meta.get("max_length", 256)
-BACKBONE: str = meta.get("backbone", "microsoft/deberta-v3-base")
-
+# ----------------------------
+# Load tokenizer
+# ----------------------------
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
-
+# ----------------------------
+# Define model class
+# ----------------------------
 class EmotionClassifier(nn.Module):
-    def __init__(self, model_name: str, num_labels: int, dropout_rate: float = 0.3):
+    def __init__(self, model_name, num_labels=5, dropout_rate=0.3):
         super().__init__()
         self.config = AutoConfig.from_pretrained(model_name)
         self.transformer = AutoModel.from_pretrained(model_name)
         self.dropout = nn.Dropout(dropout_rate)
         self.classifier = nn.Linear(self.config.hidden_size, num_labels)
 
-    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    def forward(self, input_ids, attention_mask):
         outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
         cls = outputs.last_hidden_state[:, 0, :]
         cls = self.dropout(cls)
-        return self.classifier(cls)
-
-
-def build_and_load_model(backbone: str, weights_path: str, num_labels: int) -> nn.Module:
-    model = EmotionClassifier(backbone, num_labels=num_labels)
-    device = DEVICE
-    if Path(weights_path).exists():
-        state = torch.load(weights_path, map_location=device)
-        if isinstance(state, dict) and "state_dict" in state:
-            state = state["state_dict"]
-        cleaned = {k.replace("module.", ""): v for k, v in state.items()}
-        model.load_state_dict(cleaned, strict=False)
-    model.to(device)
-    model.eval()
-    return model
-
-
-model = build_and_load_model(BACKBONE, WEIGHTS_PATH, num_labels=len(TARGET_COLS))
-
-
-def predict_text(text: str, threshold: float = 0.5) -> Tuple[Dict[str, float], Dict[str, int]]:
-    inputs = tokenizer(text, truncation=True, padding=True, max_length=MAX_LENGTH, return_tensors="pt")
-    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+        logits = self.classifier(cls)
+        return logits
+
+# ----------------------------
+# Load model
+# ----------------------------
+model = EmotionClassifier(BACKBONE, num_labels=len(TARGET_COLS))
+state = torch.load(WEIGHTS_PATH, map_location=DEVICE)
+# Strip 'module.' prefix if present
+new_state = {k.replace("module.", ""): v for k,v in state.items()}
+model.load_state_dict(new_state, strict=False)
+model.to(DEVICE)
+model.eval()
+
+# ----------------------------
+# Prediction function
+# ----------------------------
+def predict_text(text, threshold=0.5):
+    inputs = tokenizer(
+        text,
+        truncation=True,
+        padding=True,
+        max_length=MAX_LENGTH,
+        return_tensors="pt"
+    )
+    inputs = {k:v.to(DEVICE) for k,v in inputs.items()}
+    
     with torch.no_grad():
         logits = model(**inputs)
-    probs = torch.sigmoid(logits).cpu().numpy().ravel()
-    percentages = {label: round(float(p) * 100.0, 2) for label, p in zip(TARGET_COLS, probs)}
-    binary = {label: int(p > threshold) for label, p in zip(TARGET_COLS, probs)}
+    
+    probs = torch.sigmoid(logits).cpu().numpy()[0]
+    
+    # Percentages
+    percentages = {label: round(float(p)*100,2) for label,p in zip(TARGET_COLS, probs)}
+    
+    # Binary predictions
+    binary = {label: int(p>threshold) for label,p in zip(TARGET_COLS, probs)}
+    
+    # Sort probabilities descending for display
     percentages_sorted = dict(sorted(percentages.items(), key=lambda x: x[1], reverse=True))
+    
     return percentages_sorted, binary
 
-
-def create_gradio_app():
-    demo = gr.Blocks()
-    with demo:
-        gr.Markdown("# Multi-label Emotion Classifier")
-        txt = gr.Textbox(lines=4, placeholder="Enter text here...", label="Input Text")
-        thresh = gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.01, label="Threshold")
-        btn = gr.Button("Predict")
-        probs_out = gr.Label(label="Emotion Probabilities (%)")
-        binary_out = gr.JSON(label="Binary Predictions (0/1)")
-        btn.click(predict_text, inputs=[txt, thresh], outputs=[probs_out, binary_out])
-    return demo
-
+# ----------------------------
+# Gradio Interface
+# ----------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("# Multi-label Emotion Classifier")
+    txt = gr.Textbox(lines=4, placeholder="Enter text here...", label="Input Text")
+    thresh = gr.Slider(minimum=0.1, maximum=0.9, value=0.5, step=0.01, label="Threshold")
+    btn = gr.Button("Predict")
+    probs_out = gr.Label(label="Emotion Probabilities (%)")
+    binary_out = gr.JSON(label="Binary Predictions (0/1)")
+    
+    btn.click(predict_text, inputs=[txt, thresh], outputs=[probs_out, binary_out])
 
 if __name__ == "__main__":
-    app = create_gradio_app()
-    app.launch()
+    demo.launch()