import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

# 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:
    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):
    if classifier is None:
        return {"Error": "Model not loaded properly."}

    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(
    fn=predict,
    inputs=gr.Textbox(lines=2, placeholder="Enter a math problem..."),
    outputs="json",
    title="Sangapac Math Model",
    description="A model to classify math problems into categories like Arithmetic, Multiplication, Division, Algebra, and Geometry.",
)

# Launch the app
interface.launch()