File size: 4,008 Bytes
2405121 | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Ryskina & Knight (2021) next-term eval for byte-level spark-gpt OEIS checkpoints.
Each item is a short sequence. The prompt is every term but the last, written
"t1,t2,...,tk," with a trailing comma. The answer is the last term.
Two scores per item:
exact greedy decode until the next comma or EOS; correct iff the text equals
str(answer). This matches the NextTerm model-card protocol.
logprob log P(answer is the next term) in nats: the sum of the log-probs of the
answer's bytes, plus log(P(",") + P(EOS)) at the position after it. A
continuous score, much less noisy than 57 binary outcomes.
Encoding is BOS + UTF-8 bytes, exactly as spark-gpt trains. The script checks
that the exported HF tokenizer agrees.
Usage: python eval_ryskina.py <hf_dir> <ryskina.jsonl> <out_prefix>
Writes <out_prefix>.items.jsonl and <out_prefix>.summary.json.
"""
import json
import math
import statistics
import sys
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MAX_NEW_TOKENS = 40 # longest answer is far shorter; a runaway decode just scores wrong
def main():
hf_dir, data_path, out_prefix = sys.argv[1:4]
tok = AutoTokenizer.from_pretrained(hf_dir)
model = AutoModelForCausalLM.from_pretrained(hf_dir, dtype=torch.float32).cuda().eval()
bos, eos = tok.bos_token_id, tok.eos_token_id
comma = ord(",")
def encode(text):
ids = [bos] + list(text.encode("utf-8"))
assert tok(text).input_ids == ids, f"tokenizer disagrees on {text!r}"
return ids
items = [json.loads(line)["seq"] for line in open(data_path)]
results = []
with torch.no_grad():
for idx, seq in enumerate(items):
prompt_text = ",".join(str(x) for x in seq[:-1]) + ","
answer = str(seq[-1])
prompt = encode(prompt_text)
# exact match: greedy until comma / EOS
ids = torch.tensor([prompt], device="cuda")
generated = []
for _ in range(MAX_NEW_TOKENS):
nxt = int(model(ids).logits[0, -1].argmax())
if nxt in (comma, eos):
break
generated.append(nxt)
ids = torch.cat([ids, torch.tensor([[nxt]], device="cuda")], dim=1)
prediction = bytes(b for b in generated if b < 256).decode("utf-8", "replace")
# log P(answer is the next term)
ans_ids = list(answer.encode("utf-8"))
full = torch.tensor([prompt + ans_ids], device="cuda")
logp = torch.log_softmax(model(full).logits[0].float(), dim=-1)
start = len(prompt) - 1 # logits at position i predict token i + 1
lp_answer = sum(float(logp[start + j, t]) for j, t in enumerate(ans_ids))
end = logp[start + len(ans_ids)]
lp_stop = float(torch.logaddexp(end[comma], end[eos]))
lp = lp_answer + lp_stop
results.append({
"idx": idx,
"prompt": prompt_text,
"answer": answer,
"prediction": prediction,
"correct": prediction == answer,
"logprob": lp,
"logprob_answer_bytes": lp_answer,
"logprob_stop": lp_stop,
})
lps = [r["logprob"] for r in results]
n_correct = sum(r["correct"] for r in results)
summary = {
"checkpoint": hf_dir,
"items": len(results),
"exact_correct": n_correct,
"exact_acc": n_correct / len(results),
"mean_logprob": statistics.fmean(lps),
"median_logprob": statistics.median(lps),
"mean_answer_prob": statistics.fmean(math.exp(x) for x in lps),
}
with open(out_prefix + ".items.jsonl", "w") as f:
for r in results:
f.write(json.dumps(r) + "\n")
with open(out_prefix + ".summary.json", "w") as f:
json.dump(summary, f, indent=2)
print(json.dumps(summary, indent=2))
if __name__ == "__main__":
main()
|