hfx_runtime.py

"""
Runtime pieces shared by inference and evaluation.

Everything needed to LOAD and RUN a released checkpoint lives here: the
projector architecture, the sentinel that marks where soft tokens are spliced
in, and the dataset loader. Training code is intentionally not part of this
module — the released artefact is the checkpoint, not the training loop.
"""
from __future__ import annotations
import json

import torch
import torch.nn as nn
import torch.nn.functional as F

# Marks the splice point in the prompt where soft tokens replace text.
SENTINEL = "<<<MRI_SOFT_TOKENS>>>"


class Projector(nn.Module):
    """
    Classifier representations → k soft tokens in the LLM's embedding space.

    Follows Gemma 4 Unified's patch path: LayerNorm → Dense → LayerNorm, then
    the shared multimodal embedder pattern (RMSNorm → Linear). The learnable
    token-position embedding lets the k tokens differentiate from one another —
    a flattened version of Gemma's factorized 2D positional embedding, since our
    sequence is linear rather than a grid.
    """

    def __init__(self, in_dim: int, hidden: int, n_tokens: int = 4, mid: int = 2048,
                 target_norm: float | None = None):
        """
        target_norm: the mean L2 norm of the LLM's own token embeddings. When
        given, the output is rescaled to it.

        Why: left unconstrained, the projector emits vectors at a magnitude the
        model has never seen, and generation degenerates into looping the same
        phrase. That is exactly what the first run did — output with soft tokens
        was broken while the zeroed-out version was fine. Fixing the norm
        prevents it, and is standard practice in VLM training.
        """
        super().__init__()
        self.n_tokens = n_tokens
        self.hidden = hidden
        self.pre = nn.Sequential(
            nn.LayerNorm(in_dim),
            nn.Linear(in_dim, mid),
            nn.GELU(),
            nn.LayerNorm(mid),
        )
        self.to_tokens = nn.Linear(mid, n_tokens * hidden)
        self.pos = nn.Parameter(torch.zeros(1, n_tokens, hidden))
        self.out_norm = nn.RMSNorm(hidden) if hasattr(nn, "RMSNorm") else nn.LayerNorm(hidden)
        self.out = nn.Linear(hidden, hidden)
        nn.init.normal_(self.pos, std=0.02)
        self.register_buffer("target_norm",
                             torch.tensor(float(target_norm)) if target_norm else
                             torch.tensor(0.0))

    def forward(self, feats: torch.Tensor) -> torch.Tensor:      # (B, in_dim)
        h = self.pre(feats)
        t = self.to_tokens(h).view(-1, self.n_tokens, self.hidden)
        t = t + self.pos
        t = self.out(self.out_norm(t))
        if float(self.target_norm) > 0:
            t = F.normalize(t, dim=-1) * self.target_norm
        return t                                                  # (B, k, hidden)


def apply_template(tok, prompt_text: str) -> str:
    """Chat template with thinking OFF (targets contain no chain of thought)."""
    msgs = [{"role": "user", "content": prompt_text}]
    kw = dict(tokenize=False, add_generation_prompt=True)
    try:
        return tok.apply_chat_template(msgs, enable_thinking=False, **kw)
    except TypeError:
        return tok.apply_chat_template(msgs, **kw)


def build_examples(features_path: str, text_path: str):
    """Join the cached classifier features with the rendered text dataset."""
    d = torch.load(features_path, map_location="cpu", weights_only=False)
    with open(text_path, encoding="utf-8") as f:
        td = json.load(f)

    feats = torch.cat([d["fused_features"], d["mri_features"], d["tab_features"]], dim=-1)
    recs = td["records"]
    # The multitask file holds several records per patient sharing one index, so
    # row counts are not expected to match — only the index bound is checked.
    max_idx = max(r["index"] for r in recs)
    if max_idx >= feats.size(0):
        raise ValueError(f"index {max_idx} in the text file exceeds the feature "
                         f"cache size ({feats.size(0)}) — same cache?")

    cls = list(d["class_names"])
    head_cls = [cls[int(i)] for i in d["class_probs"].argmax(dim=-1)]

    out = {"train": [], "val": [], "test": []}
    for r in recs:
        out[r["split"]].append({
            "feat": feats[r["index"]],
            "prompt": r["prompt"],
            "target": r["target"],
            "label": r["label"],
            "ptid": r["ptid"],
            # The head's verdict is read DIRECTLY from class_probs. Evaluation
            # used to regex it out of the prompt, which produced head=None
            # whenever the prompt omitted that line and silently broke the
            # faithfulness metric.
            "head": head_cls[r["index"]],
            "task": r.get("task"),
            # index is kept so chat.py can read head outputs (class_probs,
            # will_progress) from the cache and place them in the system prompt.
            "index": r["index"],
        })
    print(f"[data] train {len(out['train'])} / val {len(out['val'])} / "
          f"test {len(out['test'])} | feature dim {feats.size(1)}")
    return out, d