spb/wpai0001_uqo

1"""Prompt-level helpers for Tier-1 experiments on mlx_lm models: last-token hidden states per layer,2greedy generation with an intervention applied during the prompt pass only, and next-token NLL."""3from __future__ import annotations45import numpy as np67import mlx.core as mx8from mlx_lm.models.cache import make_prompt_cache910from wp_ai_0001.models.mlx_backend import MLXCapture, instrument111213class Runner:14    def __init__(self, model, tokenizer, name: str) -> None:15        self.model, self.tok, self.name = model, tokenizer, name16        self.hooks = instrument(model)17        self.n_layers = len(model.model.layers)18        self.d = model.args.hidden_size1920    def encode(self, text: str) -> list[int]:21        return self.tok.encode(text)2223    def last_hidden(self, text: str, layers: list[int]) -> dict[int, np.ndarray]:24        cap = MLXCapture("last")25        names = [f"resid_post.{l}" for l in layers]26        for n in names:27            self.hooks.add(n, cap)28        try:29            self.model(mx.array([self.encode(text)]))30            mx.eval(*[cap.store[n] for n in names])31        finally:32            for n in names:33                self.hooks.remove(n, cap)34        return {l: np.array(cap.store[f"resid_post.{l}"][0].astype(mx.float32)) for l in layers}3536    def generate(self, text: str, max_tokens: int = 6, intervention=None, layer: int | None = None) -> str:37        """Greedy decoding; ``intervention`` (callable(x, ctx)) is active only during the prompt pass at38        ``resid_post.{layer}``; subsequent steps run with hooks disabled (within-sequence persistence via KV)."""39        ids = self.encode(text)40        cache = make_prompt_cache(self.model)41        if intervention is not None:42            self.hooks.add(f"resid_post.{layer}", intervention)43        try:44            logits = self.model(mx.array([ids]), cache=cache)45        finally:46            if intervention is not None:47                self.hooks.remove(f"resid_post.{layer}", intervention)48        out = []49        prev = self.hooks.enabled50        self.hooks.enabled = False51        try:52            for _ in range(max_tokens):53                tok = int(mx.argmax(logits[0, -1, :]).item())54                out.append(tok)55                piece = self.tok.decode(out)56                if "\n" in piece:57                    break58                logits = self.model(mx.array([[tok]]), cache=cache)59        finally:60            self.hooks.enabled = prev61        return self.tok.decode(out)6263    def next_token_nll(self, text: str, intervention=None, layer: int | None = None) -> float:64        """NLL of the last token of ``text`` given the preceding tokens, with an optional intervention at the65        prompt's last position (i.e. the position predicting that token)."""66        ids = self.encode(text)67        ctx, target = ids[:-1], ids[-1]68        if intervention is not None:69            self.hooks.add(f"resid_post.{layer}", intervention)70        try:71            logits = self.model(mx.array([ctx]))[0, -1, :].astype(mx.float32)72        finally:73            if intervention is not None:74                self.hooks.remove(f"resid_post.{layer}", intervention)75        lp = logits - mx.logsumexp(logits)76        return float(-lp[target].item())77