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