Update app.py

app.py

@@ -3,13 +3,6 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassifica
 
 # Model details
 MODEL_NAME = "Pisethan/sangapac-math"
-reverse_label_mapping = {
-    0: "arithmetic",
-    1: "multiplication",
-    2: "division",
-    3: "algebra",
-    4: "geometry",
-}
 
 # Load model and tokenizer
 try:
@@ -27,12 +20,12 @@ def predict(input_text):
     try:
         # Predict the category
         result = classifier(input_text)
-        label_id = int(result[0]["label"].split("_")[-1])  # Extract label ID
-        category = reverse_label_mapping[label_id]  # Map label to category
+        label = result[0]["label"]  # Directly use the label provided by the model
+        score = result[0]["score"]
 
         return {
-            "Category": category,
-            "Confidence": result[0]["score"],
+            "Category": label,  # No need to remap since the label is already readable
+            "Confidence": score,
         }
     except Exception as e:
         return {"Error": str(e)}