File size: 1,461 Bytes
ea7e326 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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()
|