Update app.py

app.py

@@ -6,7 +6,7 @@ MODEL_NAME = "blaikhole/distilbert-review-bug-classifier"
 
 # Load model config to get label mapping
 config = AutoConfig.from_pretrained(MODEL_NAME)
-id2label = config.id2label 
+id2label = config.id2label  # Example: {0: "Bug", 1: "Feature Request", 2: "Documentation"}
 
 # Create a pipeline for text classification
 pipe = pipeline("text-classification", model=MODEL_NAME)
@@ -22,14 +22,13 @@ user_input = st.text_area("Input Text:", height=150)
 if st.button("Classify"):
     if user_input:
         result = pipe(user_input, return_all_scores=True)[0]  # Get all scores
-        
+
         # Convert "LABEL_n" to actual class names
         predictions = {id2label[int(res['label'].replace('LABEL_', ''))]: res['score'] for res in result}
         top_label = max(predictions, key=predictions.get)
 
-        # Display results
-        st.write(f"### 🏆 Predicted Category: `{top_label}`")
+        # Show decoded label
+        st.write(f"### 🏆 Predicted Category: `{top_label}`})")
         st.json(predictions)  # Show confidence scores
     else:
         st.warning("⚠️ Please enter some text.")
-