Update app.py

app.py

@@ -9,17 +9,26 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Access label mapping directly from model config
+# Model's id2label from config
 id2label = model.config.id2label
 
+# Custom label mapping
+label_map = {
+    "label_0": "Graphical Issue",
+    "label_1": "Network Issue",
+    "label_2": "No Bug ✅",
+    "label_3": "Performance Issue"
+}
+
 # Classification function
 def classify_review(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128).to(device)
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_label_id = torch.argmax(outputs.logits).item()
-    predicted_label = id2label[predicted_label_id]
-    return predicted_label
+    hf_label = id2label[predicted_label_id]
+    custom_label = label_map.get(hf_label, "Unknown")
+    return custom_label
 
 # Gradio interface
 iface = gr.Interface(