Natural Questions

pip install -U datasets transformers evaluate lm-eval sacrebleu rouge-score scikit-learn

from datasets import load_dataset

ds = load_dataset("google-research-datasets/nq_open", split="validation")
print(ds)            # columns: question, answer (list of acceptable strings)
print(ds[0])

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL = "gpt2"                                   # or "Qwen/Qwen2.5-1.5B", "meta-llama/Llama-3.2-1B", …
tok = AutoTokenizer.from_pretrained(MODEL); model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16).eval().to("cuda")

@torch.no_grad()
def loglik(context, continuation):
    """Sum of log-probs of `continuation` tokens given `context` — the standard multiple-choice scorer."""
    ids = tok(context + continuation, return_tensors="pt").input_ids.to(model.device)
    n_ctx = len(tok(context).input_ids)
    logp = model(ids).logits.log_softmax(-1)[0, :-1]
    return logp.gather(1, ids[0, 1:, None])[n_ctx - 1:].sum().item()

@torch.no_grad()
def generate(prompt, max_new_tokens=32, stop=None):
    ids = tok(prompt, return_tensors="pt").input_ids.to(model.device)
    out = model.generate(ids, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tok.eos_token_id)
    text = tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
    return text.split(stop)[0] if stop else text

import re, string
def norm(s):
    s = s.lower(); s = "".join(c for c in s if c not in string.punctuation)
    s = re.sub(r"\b(a|an|the)\b", " ", s); return " ".join(s.split())
FEW = "Q: who wrote the novel 1984?\nA: George Orwell\n\nQ: what is the capital of australia?\nA: Canberra\n\n"
em = 0
for ex in ds:
    golds = ex.get("answer") or ex.get("answers") or ex["answer"]["aliases"]
    pred = generate(FEW + f"Q: {ex['question']}\nA:", max_new_tokens=16, stop="\n")
    em += int(norm(pred) in {norm(g) for g in golds})
print("exact match", em / len(ds))

lm_eval --model hf --model_args pretrained=gpt2,dtype=bfloat16 \
        --tasks nq_open --num_fewshot 0 --batch_size auto --output_path results/