File size: 2,939 Bytes
84bcfb1 | 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 | """Find real reasoning loops in finished runs, down to the token where the loop starts."""
from __future__ import annotations
import glob
import gzip
import json
import os
import numpy as np
from .protocol import render
def find_loop(ids, max_period: int = 1500, min_span: int = 64, min_reps: int = 3):
"""If the token list ends in one block of `period` tokens repeated back to back (at least `min_reps` times
and over at least `min_span` tokens), return (onset, period), where onset is the first token of the
repeating stretch. Otherwise None. The last token is ignored: generation can stop in the middle of a block."""
a = np.asarray(ids[:-1])
n = len(a)
for p in range(1, max_period + 1):
reps = max(min_reps, -(-min_span // p))
if reps * p > n:
break
if a[n - 1] != a[n - 1 - p]:
continue
if all(np.array_equal(a[n - (j + 1) * p:n - j * p], a[n - p:]) for j in range(1, reps)):
t = n - p - 1
while t >= 0 and a[t] == a[t + p]:
t -= 1
return t + 1, p
return None
def read_rows(lab: str, data: str) -> dict:
with open(os.path.join(lab, "data", data + ".jsonl")) as f:
return {r["id"]: r for r in map(json.loads, f)}
def iter_samples(run_dir: str):
"""Yield (record, j, sample, text) for every sample of a run (finished shards only)."""
texts = {}
for f in sorted(glob.glob(os.path.join(run_dir, "shard*.texts.jsonl.gz"))):
with gzip.open(f, "rt") as fh:
for d in map(json.loads, fh):
texts[(d["id"], d["j"])] = d["text"]
for f in sorted(glob.glob(os.path.join(run_dir, "shard*of*.jsonl"))):
with open(f) as fh:
for rec in map(json.loads, fh):
for j, s in enumerate(rec["samples"]):
if (rec["id"], j) in texts:
yield rec, j, s, texts[(rec["id"], j)]
def looping_samples(tok, run_dir: str, rows: dict, thinking: bool = True, limit: int | None = None):
"""Every sample of a run that ends in a verbatim token loop: the token ids of prompt + answer (the answer is
re-tokenized from its text, which almost always gives back the ids the model generated) and where the loop
starts, counted from the start of the prompt."""
out = []
for rec, j, s, text in iter_samples(run_dir):
if not s.get("cut"):
continue
prompt = tok(render(tok, rows[rec["id"]], thinking), add_special_tokens=False).input_ids
ans = tok(text, add_special_tokens=False).input_ids
hit = find_loop(ans)
if hit is None:
continue
onset, period = hit
out.append({"id": rec["id"], "j": j, "ids": prompt + ans, "n_prompt": len(prompt),
"onset": len(prompt) + onset, "period": period, "text": text})
if limit and len(out) >= limit:
break
return out
|