import refrom fractions import Fractionfrom transformers import AutoModelForCausalLM, AutoTokenizerimport torch name = "lxazjk/qwen2.5-1.5b-24game-grpo"tok = AutoTokenizer.from_pretrained(name)model = AutoModelForCausalLM.from_pretrained(name, torch_dtype="auto", device_map="auto") numbers = [2, 4, 7, 7]prompt = ( "You solve arithmetic target games. Use each given number exactly once. " "Only use +, -, *, / and parentheses. Think briefly, then put the final " "expression inside <answer>...</answer>. The answer must be only an " "expression, with no equals sign, result value, example, or explanation.\n\n" f"Make 24.\nNumbers: {', '.join(map(str, numbers))}\n" "Respond in this exact format:\n<think>reasoning steps</think><answer>expression only</answer>") # sample many candidates, then verifyids = tok(prompt, return_tensors="pt").to(model.device)out = model.generate(**ids, do_sample=True, temperature=1.1, top_p=0.95, num_return_sequences=16, max_new_tokens=128)cands = tok.batch_decode(out[:, ids.input_ids.shape[1]:], skip_special_tokens=True) def verify(expr, nums, target=24): toks = re.findall(r"\d+", expr) if sorted(map(int, toks)) != sorted(nums): return False if re.fullmatch(r"[0-9+\-*/() ]+", expr) is None: return False try: # evaluate with exact rational arithmetic val = eval(expr, {"__builtins__": {}}) return Fraction(val) == target except Exception: return False for c in cands: m = re.search(r"<answer>(.*?)</answer>", c, re.S) if m and verify(m.group(1).strip(), numbers): print("solved:", m.group(1).strip()); break