import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# Load model and tokenizer
model_path = "King-8/math-messenger"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
id2label = model.config.id2label

# Classification function
def classify_math_statement(statement):
    inputs = tokenizer(statement, return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    probs = F.softmax(outputs.logits, dim=1)
    pred_id = torch.argmax(probs).item()
    label = id2label[pred_id]
    confidence = round(probs[0][pred_id].item() * 100, 2)
    return f"{label} ({confidence}%)"

# Gradio interface
iface = gr.Interface(
    fn=classify_math_statement,
    inputs=gr.Textbox(lines=3, placeholder="Enter a math-related statement..."),
    outputs="text",
    title="Math Messaging Classifier",
    description="""
This model classifies math-related statements into one of the 7 messaging recommendations from **The Math Narrative Project**:

**1. ELEVATE STUDENT AGENCY** – Focus on students’ ownership, emotions, and lived experiences in learning math.  
**2. ACKNOWLEDGE REAL-WORLD CONTEXT** – Empathize with real challenges students and families face.  
**3. ACKNOWLEDGE EMOTIONS IN MATH LEARNING** – Normalize emotions like frustration and show how they’re part of learning.  
**4. MAKE MATH RELEVANT** – Highlight how math connects to real-life goals, careers, and futures.  
**5. AFFIRM THE VALUE OF MISTAKES** – Frame mistakes as essential to growth and learning.  
**6. ENCOURAGE HELP-SEEKING** – Promote the confidence and strategies needed to ask for help.  
**7. REFRAME STRUGGLE AND CAPABILITY** – Show that struggle means you’re growing, not failing.

🔗 [Learn more at The Math Narrative Project](https://www.mathnarrative.org)
""",
    examples=[
        ["Math is tough for a lot of people, and that's totally normal."],
        ["Struggling with a problem just means I'm growing my skills."],
        ["Asking for help helps me get unstuck and keep going."],
        ["I like math because it lets me solve real-life problems."],
        ["I feel empowered when I solve something on my own."],
        ["It's okay to feel confused — that's part of learning."],
        ["I don’t always get it at first, but I try again."]
    ]
)

iface.launch()