Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
-import joblib  # Assuming label_encoder is saved as a .pkl file
 
 # Load model and tokenizer
 model_name = "mmuzamilai/distilbert-review-bug-classifier"
@@ -10,17 +9,17 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# Load label encoder
-label_encoder = joblib.load("label_encoder.pkl")  # Adjust if you have another format
+# Access label mapping directly from model config
+id2label = model.config.id2label
 
 # 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 = torch.argmax(outputs.logits).item()
-    decoded_label = label_encoder.inverse_transform([predicted_label])[0]
-    return decoded_label
+        predicted_label_id = torch.argmax(outputs.logits).item()
+    predicted_label = id2label[predicted_label_id]
+    return predicted_label
 
 # Gradio interface
 iface = gr.Interface(