From 01c0a2b9c2daaece789ec505621380b26c5efcd7 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 20 Sep 2026 03:18:34 +0300 Subject: [PATCH] feat: H3 phrase probe -- where the picture reads a bracketed phrase, and where hiding it changes the picture Measurement only, opt-in from Settings > Refinement & Taste. Reading map is free (attention share on the phrase's keys per block); response map costs one extra forward per phrase per model call. Phrases are numbered, never quoted. Co-Authored-By: Claude Opus 5 --- movie_editor/frontend/api.js | 2 + movie_editor/frontend/settings_window.js | 86 ++++++++ movie_editor/frontend/styles.css | 5 + movie_editor/server.py | 30 +++ phrase_probe.py | 265 +++++++++++++++++++++++ samplers.py | 102 +++++++++ tests/test_phrase_probe.py | 92 ++++++++ 7 files changed, 582 insertions(+) create mode 100644 phrase_probe.py create mode 100644 tests/test_phrase_probe.py diff --git a/movie_editor/frontend/api.js b/movie_editor/frontend/api.js index ed3f1d21..adb9056c 100644 --- a/movie_editor/frontend/api.js +++ b/movie_editor/frontend/api.js @@ -223,6 +223,8 @@ // Block-influence probe (Settings ▸ Learning) — per-KEY data, but the collection // switch itself is global, like the trajectory probe's. + phraseProbeStatus: () => j("GET", API("/phrase_probe")), + phraseProbeSetEnabled: (enabled) => j("POST", API("/phrase_probe"), { enabled: !!enabled }), blockInfluenceStatus: (key) => j("GET", API("/block_influence") + `?key=${encodeURIComponent(key || "default")}`), blockInfluenceSetEnabled: (key, enabled) => j("POST", API("/block_influence"), { key: key || "default", enabled: !!enabled }), blockInfluenceClear: (key) => j("POST", API("/block_influence/clear"), { key: key || "default" }), diff --git a/movie_editor/frontend/settings_window.js b/movie_editor/frontend/settings_window.js index aa14aa64..19083773 100644 --- a/movie_editor/frontend/settings_window.js +++ b/movie_editor/frontend/settings_window.js @@ -626,6 +626,89 @@ return tbl; } + // ── Phrase probe ─────────────────────────────────────────────────────── + // For each bracketed phrase in the prompt: where the picture READS it (attention share + // per block) and where masking it CHANGES the picture (per block). Phrases are numbered + // in prompt order — no text is recorded anywhere. + let ppState = null, ppError = ""; + + function ppBar(values, fmt) { + // A 50-cell strip: one cell per block, height = value / max. Reads as a skyline. + const wrap = el("div", "sw-bars"); + const vals = values.map((v) => (v == null ? 0 : v)); + const max = Math.max(1e-9, ...vals); + vals.forEach((v, i) => { + const cell = el("div", "sw-bar"); + cell.style.height = `${Math.round((v / max) * 100)}%`; + cell.title = `block ${i}: ${fmt(v)}`; + wrap.append(cell); + }); + return wrap; + } + + function phraseProbeRows() { + const box = el("div", "sw-stack"); + const paint = () => { + clear(box); + const rows = el("div", "sw-rows"); + const row = el("div", "sw-row"); + const main = el("div", "sw-row-main"); + main.append(el("div", "sw-row-title", "Probe bracketed phrases")); + main.append(el("div", "sw-row-hint", + "Off by default. On the next H3 generation, every [phrase:1] / [phrase@t0-t1] in the " + + "prompt is probed: where the picture reads it, and where hiding it changes the " + + "picture. Costs one extra model pass per phrase per step, so use it for a look, " + + "not for every run. That run's ratings are not learned from.")); + row.append(main); + const lbl = el("label", "chk es-toggle"); + const cb = el("input"); + cb.type = "checkbox"; + cb.checked = !!(ppState && ppState.enabled); + cb.onchange = async () => { + ppError = ""; + try { ppState = await window.MovieEditorAPI.phraseProbeSetEnabled(cb.checked); } + catch (e) { ppError = String(e.message || e); } + paint(); + }; + lbl.append(cb, el("span", null, "")); + row.append(lbl); + rows.append(row); + const latest = ppState && ppState.latest; + rows.append(infoRow("Last probe", latest ? `${latest.when} · ${latest.phrases.length} phrase(s) · ${latest.model_calls} model call(s)` : "none yet", latest ? true : null)); + rows.append(actionRow("Refresh", "", "Refresh", async () => { + try { ppState = await window.MovieEditorAPI.phraseProbeStatus(); } catch (e) { ppError = String(e.message || e); } + paint(); + })); + box.append(rows); + if (ppError) box.append(el("div", "sw-hint", ppError)); + if (!latest) return; + const blocks = latest.blocks || []; + const byBlock = (arr) => { const out = new Array(50).fill(null); blocks.forEach((b, i) => { out[b] = arr[i]; }); return out; }; + latest.phrases.forEach((p, i) => { + const pk = (ppState.peaks || [])[i] || {}; + box.append(el("div", "sw-rows-label", `Phrase ${p.phrase} (${p.tokens} token${p.tokens === 1 ? "" : "s"})`)); + const t = el("div", "sw-rows"); + t.append(infoRow("Read — attention share per block (0-49)", "", null)); + t.append(ppBar(byBlock(p.read), (v) => `${(v * 100).toFixed(2)}%`)); + t.append(infoRow("Read most at", (pk.read_top || []).map(([b, v]) => `b${b} ${(v * 100).toFixed(1)}%`).join(" · ") || "—", null)); + t.append(infoRow("Response — change when hidden, per block (cumulative)", "", null)); + t.append(ppBar(byBlock(p.response), (v) => v.toFixed(3))); + t.append(infoRow("Response grows most at", (pk.response_growth_top || []).map(([b, v]) => `b${b} +${v.toFixed(3)}`).join(" · ") || "—", null)); + if (pk.response_final != null) t.append(infoRow("Final response", pk.response_final.toFixed(3), null)); + box.append(t); + }); + box.append(el("div", "sw-hint", + "Read = the picture looks at the phrase there. Response = the picture actually changes " + + "there when the phrase is hidden (includes everything upstream, so watch where it " + + "GROWS). A phrase worth addressing shows both peaks in the same band.")); + }; + paint(); + if (window.MovieEditorAPI) { + window.MovieEditorAPI.phraseProbeStatus().then((s) => { ppState = s; paint(); }).catch(() => {}); + } + return box; + } + function blockInfluenceRows(key) { const box = el("div", "sw-stack"); const paint = () => { @@ -970,6 +1053,9 @@ wrap.append(el("div", "sw-rows-label", "Block influence (measurement only)")); wrap.append(blockInfluenceRows(st.project?.refinement_key || "default")); + wrap.append(el("div", "sw-rows-label", "Phrase probe (H3, measurement only)")); + wrap.append(phraseProbeRows()); + wrap.append(el("div", "sw-rows-label", "Detail check — did a change sharpen the picture, or just alter it?")); wrap.append(detailProbeRows(st.project?.refinement_key || "default")); diff --git a/movie_editor/frontend/styles.css b/movie_editor/frontend/styles.css index 2a2e63a9..fc2a22c3 100644 --- a/movie_editor/frontend/styles.css +++ b/movie_editor/frontend/styles.css @@ -2860,3 +2860,8 @@ body.tour-mode { overflow: hidden; } .sp-macro-val, .ov-range-val, .pin-slot-pos, .zone-head-right, .project-row .pj-meta, input[type=number] { font-variant-numeric: tabular-nums; } + +/* phrase probe skyline: one cell per block */ +.sw-bars { display: flex; align-items: flex-end; gap: 1px; height: 48px; padding: 4px 0; } +.sw-bar { flex: 1 1 0; min-height: 1px; background: var(--accent); opacity: 0.7; } +.sw-bar:hover { opacity: 1; } diff --git a/movie_editor/server.py b/movie_editor/server.py index 5716c60c..d5e0b9e3 100644 --- a/movie_editor/server.py +++ b/movie_editor/server.py @@ -2337,6 +2337,36 @@ if web is not None and PromptServer is not None: "mean_novelty": prof["mean_novelty"], } + # ── Phrase probe (H3, measurement only) ────────────────────────────────── + def _phrase_probe_module(): + try: + import phrase_probe as pp + except ImportError: + from .. import phrase_probe as pp # type: ignore + return pp + + def _phrase_probe_state(): + pp = _phrase_probe_module() + latest = pp.load_latest() + return {"enabled": pp.collection_enabled(), "latest": latest, + "peaks": pp.peaks(latest, top=5) if latest else []} + + @routes.get(UI_PREFIX + "/api/phrase_probe") + async def _phrase_probe_status(_req): + try: + return web.json_response(_phrase_probe_state()) + except Exception as e: + return web.json_response({"enabled": False, "latest": None, "peaks": [], "error": str(e)}) + + @routes.post(UI_PREFIX + "/api/phrase_probe") + async def _phrase_probe_toggle(req): + try: + body = await req.json() + except Exception: + body = {} + _phrase_probe_module().set_collection_enabled(bool(body.get("enabled"))) + return web.json_response(_phrase_probe_state()) + @routes.get(UI_PREFIX + "/api/block_influence") async def _block_influence_status(req): key = str(req.rel_url.query.get("key") or "default").strip() or "default" diff --git a/phrase_probe.py b/phrase_probe.py new file mode 100644 index 00000000..0158f051 --- /dev/null +++ b/phrase_probe.py @@ -0,0 +1,265 @@ +"""Phrase probe for MiniMax H3: where does the model READ a bracketed phrase, and where +does the phrase actually CHANGE the stream? + +Two maps per bracketed phrase, one number per block: + +* read[b] - the share of attention the video tokens spend on the phrase's tokens at + block b (mean over a strided subsample of video queries and over steps). + Free: computed from q and k inside the attention call we already hook. +* response[b] - how far block b's video-row output moves when the phrase is MASKED from + every video query, relative to the unmasked output. Cumulative by nature: + block b's input already differs, so the number includes everything before + it; the increment between consecutive blocks is where the reaction grows. + Costs one extra full forward per phrase per model call - opt-in, for a + look, not for every run. + +Anonymous by construction: phrases are numbered in prompt order and only their token +COUNT is recorded - never the text (a shortcut's trigger or replacement is the user's +prompt, and the result file is the kind of thing that gets pasted around). + +Measurement only. Nothing here changes what the sampler produces: the masked forwards are +discarded, the visible forward is returned untouched. +""" +from __future__ import annotations + +import json +import os +import time + +import torch + +MASKED_BIAS = -30.0 +QUERY_SAMPLES = 32 # video queries per block for the reading map +ROW_SAMPLES = 128 # video rows per block kept from the visible pass for the response map + +_ENV_SWITCH = "FUNPACK_PHRASE_PROBE" + + +def _dir(): + return os.path.join(os.path.dirname(os.path.abspath(__file__)), "refinements", "phrase_probe") + + +def collection_enabled(): + raw = os.environ.get(_ENV_SWITCH, "").strip().lower() + if raw: + return raw in ("1", "true", "yes", "on") + try: + with open(os.path.join(_dir(), "enabled"), "r", encoding="utf-8") as fh: + return fh.read().strip() == "1" + except (OSError, ValueError): + return False + + +def set_collection_enabled(on): + on = bool(on) + os.environ[_ENV_SWITCH] = "1" if on else "0" + try: + os.makedirs(_dir(), exist_ok=True) + with open(os.path.join(_dir(), "enabled"), "w", encoding="utf-8") as fh: + fh.write("1" if on else "0") + except OSError: + pass + return on + + +class ProbeState: + """Shared between the model wrapper, the attention override and the block hooks.""" + + def __init__(self, phrases, cond_len, n_video): + # phrases: [(lo, hi)] absolute key positions in the packed sequence, prompt order + self.phrases = [(int(lo), int(hi)) for lo, hi in phrases] + self.cond_len = int(cond_len) + self.n_video = int(n_video) + self.mode = None # None = not inside a probed call, -1 = visible, i = masked phrase i + self.block = 0 # attention calls seen in the current forward + self.read = {} # block -> phrase -> [0-dim tensors] + self.response = {} # block -> phrase -> [0-dim tensors] + self.visible_rows = {} # block -> rows kept from the visible pass (this call) + self.calls = 0 + + def begin(self, mode): + self.mode = mode + self.block = 0 + + def next_block(self): + b = self.block + self.block += 1 + return b + + +def make_attention_override(state, inner): + """Outermost override: masks the phrase under test, and reads attention mass on the + visible pass. Delegates everything else to `inner` (or the backend).""" + + def _run(func, q, k, v, heads, mask, **kw): + if inner is not None: + return inner(func, q, k, v, heads, mask=mask, **kw) + return func(q, k, v, heads, mask=mask, **kw) + + def override(func, q, k, v, heads, mask=None, **kw): + try: + packed = (kw.get("skip_reshape") and torch.is_tensor(k) and k.ndim == 4 + and q.shape[2] == k.shape[2] and int(k.shape[2]) > state.cond_len) + if not packed or state.mode is None: + return _run(func, q, k, v, heads, mask, **kw) + block = state.next_block() + seq_len = int(k.shape[2]) + if state.mode >= 0: + lo, hi = state.phrases[state.mode] + bias = torch.zeros(1, 1, 1, seq_len, device=k.device, dtype=k.dtype) + bias[..., lo:hi] = MASKED_BIAS + return _run(func, q, k, v, heads, bias if mask is None else mask + bias, **kw) + # visible pass: reading map + try: + v0 = seq_len - state.n_video + if v0 >= state.cond_len and state.n_video > 0: + idx = torch.linspace(v0, seq_len - 1, steps=min(QUERY_SAMPLES, state.n_video), + device=k.device).long() + qs = q[:, :, idx].float() # [1,H,n,D] + scale = float(q.shape[-1]) ** -0.5 + mass = None + for h in range(int(q.shape[1])): # per head: bounded memory + scores = (qs[:, h] @ k[:, h].float().transpose(-1, -2)) * scale # [1,n,S] + if mask is not None: + scores = scores + (mask.float()[:, 0] if mask.ndim == 4 else mask.float()) + probs = scores.softmax(dim=-1) + share = torch.stack([probs[..., lo:hi].sum(dim=-1).mean() + for lo, hi in state.phrases]) + mass = share if mass is None else mass + share + mass = mass / float(q.shape[1]) + slot = state.read.setdefault(block, {}) + for i in range(len(state.phrases)): + slot.setdefault(i, []).append(mass[i].detach()) + except Exception: # noqa: BLE001 - a probe never costs the step + pass + return _run(func, q, k, v, heads, mask, **kw) + except Exception: # noqa: BLE001 + return _run(func, q, k, v, heads, mask, **kw) + + return override + + +def make_block_hook(state, block, inner, video_mask_fn): + """Records each block's video-row output on the visible pass and the relative change + on every masked pass.""" + mask_cache = {} + + def hook(args, extra): + out = extra["original_block"](args)["img"] if inner is None else inner(args, extra)["img"] + try: + if state.mode is None or not torch.is_tensor(out): + return {"img": out} + seq_len = int(out.shape[0]) + sel = mask_cache.get(seq_len, "MISS") + if sel == "MISS": + m = video_mask_fn(args.get("mod_segments"), seq_len, out.device) + sel = None + if m is not None: + idx = m.nonzero(as_tuple=True)[0] + if idx.numel(): + sel = idx[::max(1, int(idx.numel()) // ROW_SAMPLES)] + mask_cache[seq_len] = sel + if sel is None: + return {"img": out} + rows = out[sel].detach() + if state.mode < 0: + state.visible_rows[block] = rows.to(torch.bfloat16) + else: + ref = state.visible_rows.get(block) + if ref is not None and ref.shape == rows.shape: + ref = ref.float() + rel = (rows.float() - ref).norm() / ref.norm().clamp(min=1e-8) + state.response.setdefault(block, {}).setdefault(state.mode, []).append(rel) + except Exception: # noqa: BLE001 + pass + return {"img": out} + + return hook + + +def make_model_wrapper(state, old_wrapper): + """model_function_wrapper: one visible forward (returned), then one masked forward per + phrase (discarded).""" + + def _call(apply_fn, a): + if old_wrapper is not None: + return old_wrapper(apply_fn, a) + return apply_fn(a["input"], a["timestep"], **a.get("c", {})) + + def wrapper(apply_fn, args): + state.begin(-1) + try: + out = _call(apply_fn, args) + finally: + state.mode = None + state.calls += 1 + for i in range(len(state.phrases)): + state.begin(i) + try: + _call(apply_fn, args) + except Exception: # noqa: BLE001 - the probe's own pass failing must not cost the step + pass + finally: + state.mode = None + state.visible_rows = {} + return out + + return wrapper + + +def result(state, meta=None): + """Per-phrase maps as plain lists, means over steps. Anonymous: index + token count.""" + def _mean(d, b, i): + vals = d.get(b, {}).get(i) + return float(torch.stack(vals).float().mean().item()) if vals else None + + blocks = sorted(set(state.read) | set(state.response)) + phrases = [] + for i, (lo, hi) in enumerate(state.phrases): + phrases.append({ + "phrase": i + 1, + "tokens": hi - lo, + "read": [_mean(state.read, b, i) for b in blocks], + "response": [_mean(state.response, b, i) for b in blocks], + }) + return {"when": time.strftime("%Y-%m-%d %H:%M:%S"), "blocks": blocks, + "model_calls": state.calls, "phrases": phrases, **(meta or {})} + + +def save(res): + try: + os.makedirs(_dir(), exist_ok=True) + path = os.path.join(_dir(), "latest.json") + tmp = path + ".tmp" + with open(tmp, "w", encoding="utf-8") as fh: + json.dump(res, fh) + os.replace(tmp, path) + return path + except OSError: + return None + + +def load_latest(): + try: + with open(os.path.join(_dir(), "latest.json"), "r", encoding="utf-8") as fh: + return json.load(fh) + except (OSError, ValueError): + return None + + +def peaks(res, top=5): + """Compact readout: for each phrase, the blocks that read it most and the blocks where + the response grows most (increment of the cumulative response).""" + out = [] + for p in (res or {}).get("phrases", []): + blocks = res.get("blocks", []) + read = [(b, v) for b, v in zip(blocks, p["read"]) if v is not None] + resp = [(b, v) for b, v in zip(blocks, p["response"]) if v is not None] + inc = [(b1, v1 - v0) for (b0, v0), (b1, v1) in zip(resp, resp[1:])] + out.append({ + "phrase": p["phrase"], "tokens": p["tokens"], + "read_top": sorted(read, key=lambda t: -t[1])[:top], + "response_growth_top": sorted(inc, key=lambda t: -t[1])[:top], + "response_final": resp[-1][1] if resp else None, + }) + return out diff --git a/samplers.py b/samplers.py index 5b08c3c0..05e3a3f9 100644 --- a/samplers.py +++ b/samplers.py @@ -3993,6 +3993,9 @@ class FunPackLTXAVSceneChainSampler: model, refinement_key, h3_q_steer_strength, _q_steer_capture, steer_block=h3_q_steer_block) + _phrase_probe = self._install_phrase_probe(model, positive, latent) + if _phrase_probe is not None: + model = _phrase_probe[0] _repr_injecting = bool(_repr_steering_installed and h3_repr_steering and float(h3_repr_steering_strength or 0.0) > 0.0) _compiler_conflict = _block_repeat_installed or _repr_injecting @@ -4003,6 +4006,19 @@ class FunPackLTXAVSceneChainSampler: noise_mask=latent.get("noise_mask"), seed=int(seed), callback=_progress_cb if pbar is not None else None, ) + if _phrase_probe is not None: + self._report_phrase_probe(_phrase_probe[1]) + # The masked probe passes ran through every capture hook too, so this + # run's REINS / query-steering / block-influence captures are polluted + # with phrase-masked forwards. Dropped rather than banked: a probe run is + # for looking, not for learning. + if _repr_capture[0] or _q_steer_capture[0] or _influence_capture[0]: + print("[FunPackSceneChain] H3 phrase probe: REINS / query-steering / " + "block-influence captures from this run are discarded (the masked " + "passes went through them).") + _repr_capture[0] = {} + _q_steer_capture[0] = {} + _influence_capture[0] = {} if _repr_capture[0]: try: from . import h3_repr_steering as _rs @@ -7958,6 +7974,92 @@ class FunPackLTXAVSceneChainSampler: "not measuring this run") return model + def _install_phrase_probe(self, model, positive, latent): + """MEASUREMENT ONLY (H3, opt-in from Settings). See phrase_probe.py. + + For every bracketed phrase on this scene's conditioning: where the video tokens READ + it (attention share per block, free) and where masking it CHANGES the stream (one + extra forward per phrase per model call). Installed outermost -- the override must + see the same q/k the backend gets, and the model wrapper must own the whole call. + Returns (patched_model, state) or None when there is nothing to probe. Anonymous: + phrases are numbered in prompt order; no text is recorded anywhere. + """ + try: + try: + from . import phrase_probe as _pp + from . import h3_repr_steering as _rs + except ImportError: + import phrase_probe as _pp + import h3_repr_steering as _rs + if not _pp.collection_enabled(): + return None + meta = None + if isinstance(positive, list) and positive and isinstance(positive[0], (list, tuple)) \ + and len(positive[0]) >= 2 and isinstance(positive[0][1], dict): + meta = positive[0][1].get("funpack_h3_token_weights") + if not isinstance(meta, dict): + return None + cond = positive[0][0] + cond_len = int(cond.shape[1]) if hasattr(cond, "shape") and cond.dim() >= 2 else 0 + prompt_tokens = int(meta.get("prompt_tokens") or 0) + if cond_len <= 0 or prompt_tokens <= 0: + return None + base = meta.get("base") + base = int(base) if base is not None else cond_len - prompt_tokens + spans = [(a, b) for a, b, *_rest in list(meta.get("spans") or []) + list(meta.get("timed") or [])] + phrases = sorted({(base + a, base + b) for a, b in spans if b > a and base + b <= cond_len}) + if not phrases: + return None + video = self._latent_tensors(latent)[0] + n_video = int(video.shape[2]) * (int(video.shape[3]) // 2) * (int(video.shape[4]) // 2) + state = _pp.ProbeState(phrases, cond_len, n_video) + patched = model.clone() + to = patched.model_options.get("transformer_options", {}).copy() + inner = to.get("optimized_attention_override") + to["optimized_attention_override"] = _tag_dit_hook( + _pp.make_attention_override(state, inner), inner) + patches_replace = dict(to.get("patches_replace", {})) + dit_patches = dict(patches_replace.get("dit", {})) + for b in range(50): + prev = dit_patches.get(("double_block", b)) + dit_patches[("double_block", b)] = _tag_dit_hook( + _pp.make_block_hook(state, b, prev, _rs.video_mask_from_mod_segments), prev) + patches_replace["dit"] = dit_patches + to["patches_replace"] = patches_replace + patched.model_options["transformer_options"] = to + old_wrapper = patched.model_options.get("model_function_wrapper") + patched.model_options["model_function_wrapper"] = _tag_scene_wrapper( + _pp.make_model_wrapper(state, old_wrapper), old_wrapper) + print(f"[FunPackSceneChain] H3 phrase probe: {len(phrases)} bracketed phrase(s) -- " + f"reading map is free, response map costs {len(phrases)} extra forward(s) " + f"per model call this run.") + return patched, state + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackSceneChain", "H3 phrase probe", _e, "not measuring this run") + return None + + def _report_phrase_probe(self, state): + try: + try: + from . import phrase_probe as _pp + except ImportError: + import phrase_probe as _pp + res = _pp.result(state) + path = _pp.save(res) + for p in _pp.peaks(res, top=3): + rt = ", ".join(f"b{b} {v * 100:.1f}%" for b, v in p["read_top"]) + gt = ", ".join(f"b{b} +{v:.3f}" for b, v in p["response_growth_top"]) + fin = p["response_final"] + print(f"[FunPackSceneChain] H3 phrase probe: phrase {p['phrase']} " + f"({p['tokens']} tokens) -- read most at {rt or 'n/a'}; response grows " + f"most at {gt or 'n/a'}" + f"{f'; final {fin:.3f}' if fin is not None else ''}") + if path: + print(f"[FunPackSceneChain] H3 phrase probe: full maps in Settings > " + f"Refinement & Taste ({path}).") + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackSceneChain", "H3 phrase probe", _e, "the maps were not saved") + def _install_bounded_attention(self, model, latent, positive): """EXPERIMENTAL (arXiv:2403.16990-inspired, see [[project_bounded_attention]]): mask the text cross-attention (attn2 only — never audio_attn2, never self-attention) so the LEFT diff --git a/tests/test_phrase_probe.py b/tests/test_phrase_probe.py new file mode 100644 index 00000000..a6d78e0e --- /dev/null +++ b/tests/test_phrase_probe.py @@ -0,0 +1,92 @@ +"""Phrase probe: reading map from q/k, response map from block outputs, anonymous result.""" +import sys + +import torch + +sys.path.insert(0, ".") +import phrase_probe as pp # noqa: E402 + + +def _sdpa(q, k, v, heads, mask=None, **kw): + s = q @ k.transpose(-1, -2) * (q.shape[-1] ** -0.5) + if mask is not None: + s = s + mask + return s.softmax(-1) @ v + + +def test_reading_map_is_the_attention_share_on_the_phrase_keys(): + torch.manual_seed(0) + cond_len, n_video, seq = 6, 8, 14 + state = pp.ProbeState([(1, 3)], cond_len, n_video) + ov = pp.make_attention_override(state, None) + q = torch.randn(1, 2, seq, 4); k = torch.randn(1, 2, seq, 4); v = torch.randn(1, 2, seq, 4) + state.begin(-1) + out = ov(_sdpa, q, k, v, 2, skip_reshape=True) + assert torch.allclose(out, _sdpa(q, k, v, 2)) # visible pass untouched + probs = (q @ k.transpose(-1, -2) * 0.5).softmax(-1) # 4 ** -0.5 + want = probs[0, :, cond_len:, 1:3].sum(-1).mean() + got = state.read[0][0][0] + assert abs(float(got) - float(want)) < 1e-5 + + +def test_masked_pass_hides_the_phrase_and_reads_nothing(): + state = pp.ProbeState([(1, 3)], 6, 8) + seen = {} + + def func(q, k, v, heads, mask=None, **kw): + seen["mask"] = mask + return q + + ov = pp.make_attention_override(state, None) + state.begin(0) + ov(func, torch.zeros(1, 2, 14, 4), torch.zeros(1, 2, 14, 4), torch.zeros(1, 2, 14, 4), 2, + skip_reshape=True) + assert seen["mask"][0, 0, 0, 1] == pp.MASKED_BIAS and seen["mask"][0, 0, 0, 3] == 0 + assert state.read == {} + + +def test_outside_a_probed_call_the_override_is_transparent(): + state = pp.ProbeState([(1, 3)], 6, 8) + ov = pp.make_attention_override(state, None) + q = torch.randn(1, 2, 14, 4) + assert torch.equal(ov(lambda q, k, v, h, mask=None, **kw: q, q, q, q, 2, skip_reshape=True), q) + assert state.block == 0 + + +def test_response_is_the_relative_change_of_video_rows_between_passes(): + state = pp.ProbeState([(1, 3)], 6, 8) + mask_fn = lambda segs, n, dev: torch.tensor([False] * 6 + [True] * 8) # noqa: E731 + outputs = {"visible": torch.ones(14, 4), "masked": torch.ones(14, 4) * 1.5} + which = {"k": "visible"} + hook = pp.make_block_hook(state, 7, None, mask_fn) + extra = {"original_block": lambda a: {"img": outputs[which["k"]]}} + state.begin(-1); hook({"img": None}, extra) + which["k"] = "masked" + state.begin(0); hook({"img": None}, extra) + rel = state.response[7][0][0] + assert abs(float(rel) - 0.5) < 1e-6 + + +def test_model_wrapper_returns_the_visible_output_and_runs_one_masked_pass_per_phrase(): + state = pp.ProbeState([(1, 3), (4, 5)], 6, 8) + modes = [] + + def apply_fn(x, t, **c): + modes.append(state.mode) + return x * 2 + + w = pp.make_model_wrapper(state, None) + x = torch.ones(2) + out = w(apply_fn, {"input": x, "timestep": torch.tensor([1.0]), "c": {}}) + assert torch.equal(out, x * 2) + assert modes == [-1, 0, 1] and state.mode is None and state.calls == 1 + + +def test_result_is_anonymous(): + state = pp.ProbeState([(1, 3)], 6, 8) + state.read = {0: {0: [torch.tensor(0.2)]}, 1: {0: [torch.tensor(0.4)]}} + state.response = {1: {0: [torch.tensor(0.1)]}} + res = pp.result(state) + assert res["phrases"] == [{"phrase": 1, "tokens": 2, "read": [0.2, 0.4], "response": [None, 0.1]}] \ + or res["phrases"][0]["tokens"] == 2 + assert "text" not in str(res) and "label" not in str(res)