| """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 |
|
|
|
|
| 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) |
|
|
| |
| 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") |
|
|
| |
| 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 |
| 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() |
|
|
