Update app.py

app.py

@@ -1,18 +1,27 @@
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 model = AutoModelForSequenceClassification.from_pretrained("King-8/interview_statements")
 
 classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
+label_map = {
+    "LABEL_0": "Education & Innovation",
+    "LABEL_1": "Community & Health",
+    "LABEL_2": "Family & History",
+    "LABEL_3": "Faith & Spirituality",
+    "LABEL_4": "Business & Finance"
+}
+
 def classify_statement(statement):
     prediction = classifier(statement)[0]
     
     label = prediction["label"]
     
-    return f"Predicted Category: {label}"
+    predicted_category = label_map.get(label, "Unknown Category")  # Default to "Unknown" if not found
+    
+    return f"Predicted Category: {predicted_category}"
 
 gr.Interface(
     fn=classify_statement,