Minutor/math-word-problem-demo

1import os2import spaces3import gradio as gr4import torch5from transformers import AutoModelForCausalLM, AutoTokenizer6from peft import PeftModel7 8BASE_MODEL = "meta-llama/Llama-3.2-3B-Instruct"9ADAPTER = "Minutor/adaption_math_word_problem_sub_2"10TOKEN = os.environ.get("HF_TOKEN")11 12model = None13tokenizer = None14 15def load_model():16    global model, tokenizer17    if model is not None:18        return19 20    print("Loading model...")21    base = AutoModelForCausalLM.from_pretrained(22        BASE_MODEL,23        torch_dtype=torch.bfloat16,24        device_map="cpu",25        token=TOKEN,26        trust_remote_code=True27    )28    model = PeftModel.from_pretrained(base, ADAPTER, token=TOKEN)29    model.eval()30 31    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=TOKEN)32    if tokenizer.pad_token is None:33        tokenizer.pad_token = tokenizer.eos_token34    print("Model loaded!")35 36SYSTEM_PROMPT = (37    "You are an expert mathematical reasoning assistant. "38    "Always provide a detailed step-by-step solution.\n"39    "Structure your answer like this:\n"40    "1. Understand: What is being asked? What numbers do we have?\n"41    "2. Plan: Should we work forwards or backwards?\n"42    "3. Solution: Show each calculation step with equations and explain why.\n"43    "4. Verify: Plug the answer back in to check.\n"44    "Final Answer: State the answer clearly.\n"45    "Avoid unnecessary variables. Only treat a quantity as zero if the problem explicitly says it is 0, zero, none, or empty - otherwise treat it as an unknown."46)47 48@spaces.GPU49def solve(question: str) -> str:50    if not question or not question.strip():51        return "Please enter a math word problem."52 53    load_model()54    model.to("cuda")55 56    messages = [57        {"role": "system", "content": SYSTEM_PROMPT},58        {"role": "user", "content": question.strip()}59    ]60 61    text = tokenizer.apply_chat_template(62        messages, tokenize=False, add_generation_prompt=True63    )64    inputs = tokenizer(text, return_tensors="pt").to("cuda")65 66    with torch.inference_mode():67        outputs = model.generate(68            **inputs,69            max_new_tokens=1024,70            do_sample=False,71            pad_token_id=tokenizer.eos_token_id72        )73 74    answer = tokenizer.decode(75        outputs[0][inputs["input_ids"].shape[1]:],76        skip_special_tokens=True77    )78    return answer.strip()79 80demo = gr.Interface(81    fn=solve,82    inputs=gr.Textbox(83        lines=4,84        placeholder="Enter a math word problem...",85        label="Math Word Problem"86    ),87    outputs=gr.Textbox(lines=14, label="Step-by-step Solution"),88    title="Math Word Problem Solver (Adaption LoRA)",89    description="Fine-tuned Llama-3.2-3B-Instruct • AutoScientist Challenge (Math & Code)",90    examples=[91        ["A farmer has chickens and cows. Altogether the animals have 50 heads and 140 legs. How many chickens and how many cows does the farmer have?"],92        ["The sum of three consecutive even integers is 150. What is the largest of these three integers?"],93        ["A laptop originally costs $800. It is first discounted by 20%, then an additional 8% tax is applied on the discounted price. What is the final price?"],94        ["I have some $5 notes and $10 notes. Altogether I have 18 notes and their total value is $130. How many $5 notes and how many $10 notes do I have?"]95    ]96)97 98demo.launch()