Update app.py

app.py

@@ -1,20 +1,35 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoConfig
+
+# Model Name (Replace with actual model)
+MODEL_NAME = "blaikhole/distilbert-review-bug-classifier"
+
+# Load model config to get label mapping
+config = AutoConfig.from_pretrained(MODEL_NAME)
+id2label = config.id2label 
 
 # Create a pipeline for text classification
-pipe = pipeline("text-classification", model="blaikhole/distilbert-review-bug-classifier")
+pipe = pipeline("text-classification", model=MODEL_NAME)
 
-# Streamlit app structure
-st.title("Review bug classification demo.")
-st.write("Enter some text and the model will predict bug category:")
+# Streamlit app UI
+st.title("Review Bug Classification Demo 🐞")
+st.write("Enter some text and the model will predict the bug category.")
 
-# Input text from the user
-user_input = st.text_input("Input Text:")
+# User Input
+user_input = st.text_area("Input Text:", height=150)
 
-# When the button is clicked, classify the input
+# Prediction
 if st.button("Classify"):
     if user_input:
-        result = pipe(user_input)
-        st.write(result)
+        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}`")
+        st.json(predictions)  # Show confidence scores
     else:
-        st.write("Please enter some text.")
+        st.warning("⚠️ Please enter some text.")
+