|
|
| import argparse, json |
| from lamm.data import stream_hf, batched |
| from lamm.verify import ProofAuditor |
| from lamm.data import format_sft |
| from lamm.metrics import summarize |
|
|
| def main(): |
| ap=argparse.ArgumentParser() |
| ap.add_argument("--dataset",default="nvidia/Nemotron-SFT-Math-v4") |
| ap.add_argument("--split",default="train") |
| ap.add_argument("--source",default="StackExchange-Math") |
| ap.add_argument("--batch-size",type=int,default=16) |
| ap.add_argument("--batches",type=int,default=10) |
| args=ap.parse_args() |
|
|
| auditor=ProofAuditor() |
| audits=[] |
| for bi,batch in enumerate(batched(stream_hf(args.dataset,args.split,args.source),args.batch_size),1): |
| print(f"BATCH {bi}: {len(batch)}") |
| for rec in batch: |
| assistant = "" |
| for m in reversed(rec.messages): |
| if isinstance(m,dict) and m.get("role")=="assistant": |
| assistant=str(m.get("content") or m.get("reasoning_content") or "") |
| break |
| if not assistant: |
| continue |
| import re |
| steps=[x.strip() for x in re.split(r"(?i)(?=\\bSTEP\\s*\\d+\\b)|\\n\\s*\\n",assistant) if x.strip()] |
| if len(steps)==1: |
| steps=[assistant] |
| a=auditor.audit(steps,rec.expected_answer) |
| audits.append(a) |
| if bi>=args.batches: break |
| print(json.dumps(summarize(audits),indent=2)) |
|
|
| if __name__=="__main__": |
| main() |
|
|
