Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
 
-# Load your model from Hugging Face Hub
+# Model details
 MODEL_NAME = "Pisethan/sangapac-math"
 reverse_label_mapping = {
     0: "arithmetic",
@@ -11,22 +11,31 @@ reverse_label_mapping = {
     4: "geometry",
 }
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
+# Load model and tokenizer
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
+except Exception as e:
+    classifier = None
+    print(f"Error loading model or tokenizer: {e}")
 
 def predict(input_text):
-    # Predict using the model
-    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
+    if classifier is None:
+        return {"Error": "Model not loaded properly."}
 
-    # Return prediction result
-    return {
-        "Category": category,
-        "Confidence": result[0]["score"],
-    }
+    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
+
+        return {
+            "Category": category,
+            "Confidence": result[0]["score"],
+        }
+    except Exception as e:
+        return {"Error": str(e)}
 
 # Gradio interface
 interface = gr.Interface(