File size: 5,175 Bytes
1013007 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """
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
|