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

# Load your model from Hugging Face Hub
MODEL_NAME = "Pisethan/sangapac-math"
reverse_label_mapping = {
    0: "arithmetic",
    1: "multiplication",
    2: "division",
    3: "algebra",
    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)

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

    # Return prediction result
    return {
        "Category": category,
        "Confidence": result[0]["score"],
    }

# 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()