import gradio as gr
from model_scratch import MathTransformer
from dataset_utils_scratch import CharTokenizer
import torch

# Load trained model (or new untrained)
tokenizer = CharTokenizer([""])  # initialize empty, replace with your trained tokenizer
model = MathTransformer(vocab_size=tokenizer.vocab_size).to("cpu")

def solve_question(text_input):
    model.eval()
    with torch.no_grad():
        x = torch.tensor([tokenizer.encode(text_input)], dtype=torch.long)
        out = model(x)
        pred_idx = torch.argmax(out, dim=-1).squeeze().tolist()
        answer = tokenizer.decode(pred_idx)
    return answer

with gr.Blocks() as demo:
    with gr.Row():
        text_input = gr.Textbox(label="Enter Math Question")
        solve_btn = gr.Button("Solve")
        output = gr.Textbox(label="Solution")
        solve_btn.click(solve_question, inputs=text_input, outputs=output)

demo.launch()