Upload eval_gsm8k_custom.py with huggingface_hub

eval_gsm8k_custom.py

@@ -6,7 +6,6 @@
 #     "torch>=2.0.0",
 #     "huggingface-hub",
 #     "accelerate",
-#     "bitsandbytes",
 # ]
 # ///
 
@@ -27,7 +26,7 @@ import torch
 
 tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
 m = AutoModelForCausalLM.from_pretrained(
-    MODEL, trust_remote_code=True, device_map="auto", torch_dtype=torch.bfloat16
+    MODEL, trust_remote_code=True, device_map="auto", torch_dtype=torch.float16
 )
 
 sys_msg = "You are a helpful assistant that solves math problems step by step."
@@ -41,7 +40,8 @@ for i, ex in enumerate(ds):
     ]
     inp = tok.apply_chat_template(msgs, tokenize=True, return_tensors="pt", add_generation_prompt=True)
     inp = {k: v.to(m.device) for k, v in inp.items()}
-    out = m.generate(**inp, max_new_tokens=512, do_sample=False)
+    with torch.no_grad():
+        out = m.generate(**inp, max_new_tokens=512, do_sample=False)
     gen = tok.decode(out[0][inp["input_ids"].shape[1]:], skip_special_tokens=True)
 
     gt_match = re.search(r"####\s*(-?\d[\d,]*\.?\d*)", ex["answer"])
@@ -56,7 +56,7 @@ for i, ex in enumerate(ds):
 
     if pred == gt:
         correct += 1
-    if i < 5:
+    if i < 5 or i % 50 == 0:
         print(f"Q{i+1}: GT={gt} Pred={pred} OK={pred==gt}")
 
 accuracy = correct / total * 100