python/elidedb/fdnnv2.py

"""FDNN-V2: the context head. Same speed class, different training signal.

V1's measured dead end: distilled to an appearance teacher whose own space
separates close/open at 0.60, no student can learn verbs — the target lacks
them. V2 keeps V1's skeleton (stem -> FDNN recurrent cell, 0.27 ms/frame)
and warm-starts from its weights; what changes is WHERE the gradient comes
from (docs/FDNNV2_PLAN.md, all research-grounded):

  L1 PREDICTIVE (V-JEPA-style, arXiv 2506.09985): from state h_t, predict
     the embedding of frame t+k at two horizons. The future is free
     supervision on every frame, and dynamics must enter h to predict it.
     This is the biological claim made mechanical: understanding = a scene
     model good enough to predict what happens next.
  L2 ARROW OF TIME (Wei et al., CVPR 2018): classify forward vs reversed
     clips from the final state. Open and close are time-reversals; one
     binary head separates them by construction.
  L3 VERB-FOCUSED CONTRASTIVE (arXiv 2304.06708): sigmoid contrastive
     between clip context embeddings and 7B captions, with hard negatives
     built by verb/direction swaps — the joint space queries actually live
     in. A 2-layer text adapter maps SigLIP text embeddings into the
     256-d context space (~0.1 ms per query).

The appearance head keeps its V1 distillation target so every existing
index and the SigLIP text tower continue to work unchanged.

THE GATE IS AN EVENT DETECTOR. The cell's update gate z_t measures how much
of the scene model this frame is allowed to overwrite. Sustained high gate
activity marks an event boundary — so the encoder emits event segmentation
for free, and event embeddings are pooled with NOVELTY WEIGHTS (gate
activity), not uniformly: the moment the drawer closes outweighs the seconds
it sat still (the TempMe lesson, arXiv 2409.01156, applied at pooling time).
"""
from __future__ import annotations

import json
import re
from pathlib import Path

import mlx.core as mx
import mlx.nn as nn
import numpy as np

from .fdnnvideo import EMBED_DIM, FDNNVideoEncoder

CTX_DIM = 256

# Verb/direction swaps for L3 hard negatives: shared lexicon (moved to
# lexicon.py so the QUERY path can import it without this module's mlx)
from .lexicon import derived_swaps  # noqa: E402,F401


def swap_verbs(text: str, rng, store=None) -> str | None:
    """One randomly chosen applicable swap -> a hard negative. None if no
    swap applies (caption has no directional content to invert)."""
    # Training negatives obey the no-hardwire rule too: a model taught
    # from hand-written oppositions has the hand-writing baked into its
    # weights, which is worse than a lookup table because it cannot be
    # grepped out afterwards.
    pairs = list(derived_swaps(store)) if store is not None else []
    if not pairs:
        return None
    t = " " + text.lower() + " "
    hits = []
    for a, b in pairs:
        if f" {a} " in t:
            hits.append((a, b))
        if f" {b} " in t:
            hits.append((b, a))
    if not hits:
        # clause order flip is the fallback inversion for compounds
        parts = re.split(r"\band then\b|\bthen\b|,", text)
        if len(parts) >= 2:
            return " then ".join(p.strip() for p in reversed(parts)
                                 if p.strip())
        return None
    a, b = hits[rng.integers(len(hits))]
    return re.sub(rf"\b{re.escape(a)}\b", b, text.lower(), count=1)


class TextAdapter(nn.Module):
    """SigLIP text embedding (1152) -> context space (256). Two layers,
    ~0.7M params, ~0.1 ms — the entire query-time cost of verb awareness."""

    def __init__(self, in_dim=EMBED_DIM, dim=CTX_DIM):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, 512)
        self.fc2 = nn.Linear(512, dim)

    def __call__(self, x):
        e = self.fc2(nn.silu(self.fc1(x)))
        return e * mx.rsqrt(mx.sum(e * e, axis=-1, keepdims=True) + 1e-8)


class FDNNv2(nn.Module):
    """V1 skeleton + context head + predictors + arrow-of-time head."""

    def __init__(self, base: FDNNVideoEncoder, ctx_dim=CTX_DIM,
                 horizons=(2, 10), motion_dim=64):
        super().__init__()
        self.base = base
        C = base.channels
        G = base.cfg["glimpse"]
        self.horizons = list(horizons)
        self.motion_dim = motion_dim
        # MOTION PATHWAY — the reversal-symmetry breaker. The gated cell is a
        # leaky integrator; for smooth inputs an EMA is nearly order-invariant,
        # and stage-A measured exactly that: AoT stuck at chance (0.47) while
        # dominating the loss. Frame-difference features flip SIGN under time
        # reversal, so direction becomes linearly readable. Velocity is not a
        # nicety here; it is the only anti-symmetric signal in the model.
        self.motion = nn.Linear(G, motion_dim)
        # context head reads (glimpse, state, motion)
        self.head_ctx = nn.Linear(G + C + motion_dim, ctx_dim)
        # Predictors target the model's OWN future glimpse-change (stop-grad).
        # A2 measured why not student-embedding changes: adjacent-frame true
        # change is ~0.2 in norm while the student's own error is ~0.45 — the
        # difference of noisy embeddings is noise, and the predictors sat at
        # cosine 0 for 8 epochs. Own-latent prediction is JEPA's actual
        # recipe; the appearance anchor stops the stem collapsing to constants.
        self.pred = [nn.Linear(C + motion_dim, G) for _ in horizons]
        # AoT reads the velocity SEQUENCE through a temporal conv, not the
        # mean: this corpus is reciprocal motion (arm out, arm back), so mean
        # signed velocity cancels — measured at chance twice. Order within
        # the window is the signal; a conv kernel can be asymmetric in time.
        self.aot_conv = nn.Conv1d(motion_dim, 32, 5, padding=2)
        self.head_aot = nn.Linear(32, 1)
        self.cfg = {"ctx_dim": ctx_dim, "horizons": self.horizons,
                    "motion_dim": motion_dim, "base": base.cfg}

    # ---- forward over a sequence, exposing everything the losses and the
    # event segmenter need ---------------------------------------------------
    def run(self, seq, h0=None):
        """(B,T,H,W,3) -> dict of app (B,T,1152), ctx (B,T,ctx), preds,
        gates (B,T), h_last."""
        b = self.base
        B, T = seq.shape[0], seq.shape[1]
        g = b.stem(seq.reshape(B * T, *seq.shape[2:])).reshape(B, T, -1)
        h = b.init_state(B) if h0 is None else h0
        hs, gates = [], []
        cell = b.cell
        for t in range(T):
            z = mx.sigmoid(g[:, t] @ cell.Wz + h @ cell.Uz + cell.bz)
            h = cell(g[:, t], h)
            hs.append(h)
            gates.append(mx.mean(z, axis=-1))
        hs = mx.stack(hs, axis=1)                    # (B,T,C)
        gates = mx.stack(gates, axis=1)              # (B,T)
        app = b.head_g(g) + b.head_h(hs)
        app = app * mx.rsqrt(mx.sum(app * app, axis=-1, keepdims=True) + 1e-8)
        # signed velocity of the glimpse; first frame gets zero motion.
        # NORMALISED: measured, ||dg|| is real motion (corr 0.63 with pixel
        # motion, 5x on moving frames) but only ~2% of the feature norm — fed
        # raw, it drowned next to signals 50x larger and every motion head
        # starved. Direction is unit-normalised; magnitude re-enters as a
        # bounded gain, so both the WHAT and the HOW-MUCH of motion survive.
        dg = mx.concatenate([mx.zeros_like(g[:, :1]),
                             g[:, 1:] - g[:, :-1]], axis=1)
        mag = mx.sqrt(mx.sum(dg * dg, axis=-1, keepdims=True) + 1e-8)
        mfeat = nn.silu(self.motion(dg / mag)) * mx.tanh(mag)
        cat = mx.concatenate([g, hs, mfeat], axis=-1)
        ctx = self.head_ctx(cat)
        ctx = ctx * mx.rsqrt(mx.sum(ctx * ctx, axis=-1, keepdims=True) + 1e-8)
        hm = mx.concatenate([hs, mfeat], axis=-1)
        preds = [p(hm) for p in self.pred]           # each (B,T,G)
        a = nn.silu(self.aot_conv(mfeat))            # (B,T,32) time-conv
        aot = self.head_aot(mx.mean(a, axis=1))[:, 0]
        return {"app": app, "ctx": ctx, "preds": preds, "gates": gates,
                "aot": aot, "h": h, "g": g}

    # ---- streaming embed (the write path): app + ctx + gate per frame -----
    def embed_stream_np(self, frames_u8, batch=128):
        h = self.base.init_state(1)
        apps, ctxs, gates = [], [], []
        for i in range(0, len(frames_u8), batch):
            x = mx.array(frames_u8[i:i + batch].astype(np.float32)
                         / 127.5 - 1.0)[None]
            out = self.run(x, h0=h)
            h = out["h"]
            apps.append(np.array(out["app"][0], np.float32))
            ctxs.append(np.array(out["ctx"][0], np.float32))
            gates.append(np.array(out["gates"][0], np.float32))
        return (np.concatenate(apps), np.concatenate(ctxs),
                np.concatenate(gates))


# ===========================================================================
# Event segmentation from the gate signal + novelty-weighted pooling
# ===========================================================================
def segment_events(gates, ts, min_len=6, smooth=5, thresh_pct=75.0):
    """Gate-activity peaks -> event boundaries.

    A boundary is where smoothed gate activity crosses above its own
    percentile threshold after having been below — the scene model is being
    rewritten. Percentile (not absolute) because gate scale is a trained
    quantity; per-stream calibration is free.
    Returns list of (start_idx, end_idx) covering the stream.
    """
    g = np.convolve(gates, np.ones(smooth) / smooth, mode="same")
    thr = np.percentile(g, thresh_pct)
    above = g > thr
    bounds = [0]
    for i in range(1, len(g)):
        if above[i] and not above[i - 1] and i - bounds[-1] >= min_len:
            bounds.append(i)
    bounds.append(len(g))
    return [(a, b) for a, b in zip(bounds[:-1], bounds[1:]) if b - a >= 2]


def pool_event(vecs, gates, lo, hi):
    """Novelty-weighted pool: frames weighted by gate activity, so change
    dominates stillness. Uniform mean is the verb-eraser; this is not."""
    w = gates[lo:hi] + 1e-3
    w = w / w.sum()
    v = (vecs[lo:hi] * w[:, None]).sum(0)
    return v / (np.linalg.norm(v) + 1e-8)


# ===========================================================================
# persistence
# ===========================================================================
def save_v2(model, adapter, meta, path):
    from mlx.utils import tree_flatten
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    np.savez(path / "v2.npz",
             **{k: np.array(v) for k, v in tree_flatten(model.parameters())})
    np.savez(path / "adapter.npz",
             **{k: np.array(v)
                for k, v in tree_flatten(adapter.parameters())})
    (path / "v2.json").write_text(json.dumps({**meta, "cfg": model.cfg},
                                             indent=2))


def load_v2(path):
    from mlx.utils import tree_unflatten
    path = Path(path)
    meta = json.loads((path / "v2.json").read_text())
    base = FDNNVideoEncoder(**meta["cfg"]["base"])
    model = FDNNv2(base, ctx_dim=meta["cfg"]["ctx_dim"],
                   horizons=tuple(meta["cfg"]["horizons"]),
                   motion_dim=meta["cfg"].get("motion_dim", 64))
    z = np.load(path / "v2.npz")
    model.update(tree_unflatten([(k, mx.array(z[k])) for k in z.files]))
    model.base.cell.freeze(keys=["basis_types", "mask"], recurse=False)
    adapter = TextAdapter()
    z = np.load(path / "adapter.npz")
    adapter.update(tree_unflatten([(k, mx.array(z[k])) for k in z.files]))
    mx.eval(model.parameters(), adapter.parameters())
    return model, adapter, meta