Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
 
 # Model name and Hugging Face token
@@ -7,30 +7,30 @@ MODEL_NAME = "Pisethan/sangapac-math"
 TOKEN = os.getenv("HF_API_TOKEN")
 
 if not TOKEN:
-    raise ValueError("Hugging Face API token not found. Set it as an environment variable.")
+    raise ValueError("Hugging Face API token not found. Set it as an environment variable (HF_API_TOKEN).")
 
 # Load model and tokenizer
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=TOKEN)
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, token=TOKEN)
-    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=TOKEN)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, use_auth_token=TOKEN)
+    generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 except Exception as e:
-    classifier = None
+    generator = None
     print(f"Error loading model or tokenizer: {e}")
 
 def predict(input_text):
-    if classifier is None:
+    if generator is None:
         return "Model not loaded properly.", {"Error": "Model not loaded properly."}
 
     try:
-        result = classifier(input_text)
-        label = result[0]["label"]
-        score = result[0]["score"]
+        # Generate output
+        result = generator(input_text, max_length=256, num_beams=5, early_stopping=True)
+        generated_text = result[0]["generated_text"]
 
-        simple_result = f"Category: {label}\nConfidence: {score:.2f}"
+        simple_result = f"Generated Solution:\n{generated_text}"
         detailed_result = {
-            "Category": label,
-            "Confidence": score,
+            "Input": input_text,
+            "Generated Solution": generated_text,
         }
 
         return simple_result, detailed_result
@@ -53,7 +53,10 @@ interface = gr.Interface(
         gr.JSON(label="Detailed JSON Output"),
     ],
     title="Sangapac Math Model",
-    description="A model to classify math problems into categories like Arithmetic, Multiplication, Division, Algebra, and Geometry.",
+    description=(
+        "A model that solves math problems and provides step-by-step solutions. "
+        "Examples include Arithmetic, Multiplication, Division, Algebra, and Geometry problems."
+    ),
     examples=sample_inputs,
 )