| """Long-context evals for a checkpoint: (1) per-position CE over 8192-token |
| windows — if carried state helps, late positions beat early ones; (2) passkey |
| retrieval scored by log-prob against distractor keys at several depths. |
| |
| python eval_longctx.py --run runs/c_gala_s0 --out runs/c_gala_s0/longctx.json |
| """ |
| import argparse, json, os |
|
|
| import numpy as np |
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
| from model import HOARD, HoardConfig |
|
|
| try: |
| import tiktoken |
| ENC = tiktoken.get_encoding("gpt2") |
| except Exception: |
| ENC = None |
|
|
|
|
| def load(run_dir): |
| meta = json.load(open(os.path.join(run_dir, "config.json"))) |
| cfg = HoardConfig.from_dict(meta["config"]) |
| model = HOARD(cfg) |
| model.load_weights(os.path.join(run_dir, "model.safetensors")) |
| model.eval() |
| return model |
|
|
|
|
| def position_ce(model, toks, seq=8192, n_windows=6): |
| bins = [(0, 256), (256, 1024), (1024, 2048), (2048, 4096), (4096, 8192)] |
| sums = np.zeros(len(bins)); cnts = np.zeros(len(bins)) |
| for w in range(n_windows): |
| st = w * (seq + 1) |
| x = np.asarray(toks[st: st + seq + 1]).astype(np.int32)[None] |
| if x.shape[1] < seq + 1: |
| break |
| x = mx.array(x) |
| logits = model(x[:, :-1]) |
| ce = nn.losses.cross_entropy(logits.astype(mx.float32), x[:, 1:], reduction="none")[0] |
| mx.eval(ce) |
| ce = np.array(ce) |
| for i, (a, b) in enumerate(bins): |
| sums[i] += ce[a:b].sum(); cnts[i] += b - a |
| return {f"{a}-{b}": round(float(s / c), 4) for (a, b), s, c in zip(bins, sums, cnts)} |
|
|
|
|
| def seq_logprob(model, ids, tail_ids): |
| """log P(tail | prefix) under the model, ids = prefix+tail.""" |
| x = mx.array(np.array(ids, dtype=np.int32))[None] |
| logits = model(x[:, :-1]).astype(mx.float32) |
| lp = nn.log_softmax(logits[0], axis=-1) |
| n = len(tail_ids) |
| tgt = x[0, 1:] |
| sel = mx.take_along_axis(lp, tgt[:, None], axis=-1)[:, 0] |
| mx.eval(sel) |
| return float(np.array(sel)[-n:].sum()) |
|
|
|
|
| def passkey(model, toks, seq=8192, depths=(0.1, 0.5, 0.9), n_trials=6, seed=0): |
| if ENC is None: |
| return {"error": "tiktoken unavailable"} |
| rng = np.random.default_rng(seed) |
| out = {} |
| for depth in depths: |
| correct = 0 |
| for t in range(n_trials): |
| key = int(rng.integers(10000, 99999)) |
| distractors = [int(rng.integers(10000, 99999)) for _ in range(4)] |
| needle = ENC.encode_ordinary(f" The pass key is {key}. Remember the pass key.") |
| query = ENC.encode_ordinary(" The pass key is") |
| filler_start = int(rng.integers(0, len(toks) - seq - 1)) |
| filler = list(np.asarray(toks[filler_start: filler_start + seq]).astype(int)) |
| pos = int(depth * (seq - len(needle) - len(query) - 24)) |
| ctx = filler[:pos] + needle + filler[pos: seq - len(needle) - len(query) - 16] |
| scores = [] |
| for cand in [key] + distractors: |
| tail = ENC.encode_ordinary(f" {cand}") |
| scores.append(seq_logprob(model, ctx + query + tail, tail)) |
| correct += int(np.argmax(scores) == 0) |
| out[f"depth_{depth}"] = round(correct / n_trials, 3) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--run", required=True) |
| ap.add_argument("--data", default="data/fineweb/val.bin") |
| ap.add_argument("--out", default=None) |
| ap.add_argument("--trials", type=int, default=6) |
| ap.add_argument("--depths", type=float, nargs="+", default=[0.1, 0.5, 0.9]) |
| ap.add_argument("--seqs", type=int, nargs="+", default=[8192, 32768], |
| help="context lengths for the passkey test") |
| a = ap.parse_args() |
| model = load(a.run) |
| toks = np.memmap(a.data, dtype=np.uint16, mode="r") |
| res = {"position_ce": position_ce(model, toks), |
| "passkey_acc": {str(sq): passkey(model, toks, seq=sq, depths=tuple(a.depths), |
| n_trials=a.trials) |
| for sq in a.seqs}} |
| out = a.out or os.path.join(a.run, "longctx.json") |
| json.dump(res, open(out, "w"), indent=2) |
| print(json.dumps(res)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
