Update app.py

app.py

@@ -9,15 +9,12 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Model's id2label from config
-id2label = model.config.id2label
-
-# Custom label mapping
+# Custom label mapping using integer keys
 label_map = {
-    "label_0": "Graphical Issue",
-    "label_1": "Network Issue",
-    "label_2": "No Bug ✅",
-    "label_3": "Performance Issue"
+    0: "graphi",
+    1: "netw",
+    2: "no",
+    3: "perf"
 }
 
 # Classification function
@@ -26,9 +23,7 @@ def classify_review(text):
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_label_id = torch.argmax(outputs.logits).item()
-    hf_label = id2label[predicted_label_id]
-    custom_label = label_map.get(hf_label, "Unknown")
-    return custom_label
+    return label_map.get(predicted_label_id, "Unknown")
 
 # Gradio interface
 iface = gr.Interface(