app.py

import streamlit as st
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=MODEL_NAME)

# Streamlit app UI
st.title("Review Bug Classification Demo 🐞")
st.write("Enter some text and the model will predict the bug category.")

# User Input
user_input = st.text_area("Input Text:", height=150)

# Prediction
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}`")
        st.json(predictions)  # Show confidence scores
    else:
        st.warning("⚠️ Please enter some text.")