run_zeroshot (2).py

import json, re, argparse, torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME      = "meta-llama/Llama-3.2-3B-Instruct"  # HF model (accept license + logged in)
DATASET_NAME    = "gsm8k"                              # HF dataset
DATASET_CONFIG  = "main"                               # gsm8k config
DATASET_SPLIT   = "test"                               # evaluate only on test
BATCH_SIZE      = 16
MAX_NEW_TOKENS  = 320
OUT_PATH        = "outputs/zero_shot_cot_test.jsonl"
SEED            = 42

PROMPT_HEADER = (
    "You are a careful math tutor. Solve the problem step by step.\n"
    "End with exactly one line in the format: Final Answer: <number>\n\n"
)

def build_prompt(q: str) -> str:
    return f"{PROMPT_HEADER}Problem: {q}\n"

def _norm(s: str | None):
    if s is None: return None
    s = s.strip().replace(",", "").rstrip(".")
    return s

def extract_gold(ans: str):
    # GSM8K gold ends with: "#### 42" (supports negatives & decimals)
    m = re.search(r"####\s*(-?\d+(?:\.\d+)?)", ans)
    return _norm(m.group(1)) if m else None

def extract_pred(txt: str):
    # Last occurrence; accepts spaces/negatives/decimals/commas
    m = re.findall(r"Final\s*Answer\s*:\s*([^\n\r]+)", txt, flags=re.IGNORECASE)
    return _norm(m[-1]) if m else None

def main():
    torch.manual_seed(SEED)

    # tokenizer
    tok = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"  # consistent for generation

    # model
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.float16 if torch.cuda.is_available() else None,
        device_map="auto",
    ).eval()

    # dataset
    ds = load_dataset(DATASET_NAME, DATASET_CONFIG)[DATASET_SPLIT]

    # run eval
    total, correct = 0, 0
    # ensure output dir exists
    import os
    os.makedirs(os.path.dirname(OUT_PATH), exist_ok=True)

    with open(OUT_PATH, "w", encoding="utf-8") as f:
        for i in range(0, len(ds), BATCH_SIZE):
            batch = ds[i:i+BATCH_SIZE]
            qs    = batch["question"]
            golds = [extract_gold(a) for a in batch["answer"]]

            prompts = [build_prompt(q) for q in qs]
            enc = tok(prompts, return_tensors="pt", padding=True, truncation=True).to(model.device)

            with torch.no_grad():
                out = model.generate(
                    **enc,
                    do_sample=False,                   # greedy
                    max_new_tokens=MAX_NEW_TOKENS,
                    eos_token_id=tok.eos_token_id,
                    pad_token_id=tok.eos_token_id,
                )

            for j in range(len(qs)):
                text_full = tok.decode(out[j], skip_special_tokens=True)
                pred = extract_pred(text_full)
                ok = (pred is not None) and (golds[j] is not None) and (pred == golds[j])
                correct += int(ok); total += 1
                f.write(json.dumps({
                    "idx": i + j,
                    "question": qs[j],
                    "gold": golds[j],
                    "prediction": pred,
                    "ok": ok
                }) + "\n")

    em = correct / total
    print(f"Zero-Shot CoT EM: {correct}/{total} = {em:.4f}")
    print(f"Saved predictions → {OUT_PATH}")

if __name__ == "__main__":
    main()