import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-math-7b-instruct")
model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-math-7b-instruct", torch_dtype=torch.bfloat16, device_map="cpu")


from transformers import GenerationConfig
model.generation_config = GenerationConfig.from_pretrained("deepseek-ai/deepseek-math-7b-instruct")
model.generation_config.pad_token_id = model.generation_config.eos_token_id


def solve_math_problem(questions):
    if isinstance(questions, str):  # If input is a single string
        questions = [questions]

    results = []
    for question in questions:
        messages = [{"role": "user", "content": question}]
        input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
        outputs = model.generate(input_tensor.to(model.device), max_new_tokens=100)
        result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
        results.append(result)
    return results


interface = gr.Interface(
    fn=solve_math_problem,
    inputs="text",
    outputs="text",
    title="Math Wizard",
    description="""
        Welcome to the Math Wizard!
        Ask any math question, and let the wizard guide you through the solution step-by-step.
    """,
    allow_flagging=True,
    examples=[
        ["What is the integral of x^2?"],
        ["How do I solve a quadratic equation?"]
    ],
    theme="compact"
)

interface.launch()