Initial release: Ker-VLJEPA-3B inference package

a974113 verified 7 months ago

21.7 kB

"""
Ker-VLJEPA-3B — Inference-only model.
Loads the Llama 3.2 3B backbone, applies LoRA adapters, visual encoder,
and cross-attention bridge components for CT report generation.
Requires:
  - A local copy of meta-llama/Llama-3.2-3B (user-provided)
  - Weight files in weights/ (shipped with this package)
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
from typing import Optional, List, Tuple, Dict
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from safetensors.torch import load_file


# ---------------------------------------------------------------------------
# Visual encoder components (inference-only, no training heads)
# ---------------------------------------------------------------------------

def sinusoidal_positional_encoding(n: int, d: int) -> torch.Tensor:
    pe = torch.zeros(n, d)
    pos = torch.arange(n, dtype=torch.float).unsqueeze(1)
    div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
    pe[:, 0::2] = torch.sin(pos * div)
    pe[:, 1::2] = torch.cos(pos * div)
    return pe.unsqueeze(0)


class ZonedCrossAttention(nn.Module):
    """Z-Zoned cross-attention: each zone's queries attend only to their spatial slice range."""

    def __init__(self, slice_dim=1024, hidden_dim=1024, num_zones=32,
                 tokens_per_zone=1, num_heads=16, dropout=0.1):
        super().__init__()
        self.num_zones = num_zones
        self.tokens_per_zone = tokens_per_zone
        self.num_regions = num_zones * tokens_per_zone
        self.num_heads = num_heads
        self.hidden_dim = hidden_dim

        self.zone_queries = nn.Parameter(torch.zeros(num_zones, tokens_per_zone, hidden_dim))
        self.zone_embed = nn.Embedding(num_zones, hidden_dim)
        self.slice_proj = nn.Linear(slice_dim, hidden_dim)
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True, dropout=dropout)
        self.output_proj = nn.Linear(hidden_dim, slice_dim)
        self.fallback_pos_embed = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02)

        # Physical Z-positional encoding buffers
        div_term = torch.exp(torch.arange(0, hidden_dim, 2).float() * (-math.log(10000.0) / hidden_dim))
        self.register_buffer("div_term", div_term)

    def _add_z_pos(self, x_proj, mask):
        """Sinusoidal positional encoding based on slice index."""
        B, S, D = x_proj.shape
        device, dtype = x_proj.device, x_proj.dtype
        z_positions = torch.zeros(B, S, device=device)
        for b in range(B):
            n = int(mask[b].sum().item()) if mask is not None else S
            pos = torch.arange(n, device=device, dtype=torch.float) * 2.5 / 600.0 * 100.0
            z_positions[b, :n] = pos
        pos = z_positions.unsqueeze(-1)
        pe = torch.zeros(B, S, D, device=device, dtype=dtype)
        pe[:, :, 0::2] = torch.sin(pos * self.div_term)
        pe[:, :, 1::2] = torch.cos(pos * self.div_term)
        if mask is not None:
            pe = pe * mask.unsqueeze(-1)
        return x_proj + pe

    def forward(self, x, mask=None, metadata=None):
        B, S, _ = x.shape
        x_proj = self.slice_proj(x)
        if mask is not None:
            x_proj = self._add_z_pos(x_proj, mask)
        else:
            x_proj = x_proj + self.fallback_pos_embed

        queries = self.zone_queries.reshape(self.num_regions, self.hidden_dim)
        zone_ids = torch.arange(self.num_zones, device=x.device)
        zone_emb = self.zone_embed(zone_ids).unsqueeze(1).expand(-1, self.tokens_per_zone, -1)
        queries = (queries + zone_emb.reshape(self.num_regions, self.hidden_dim)).unsqueeze(0).expand(B, -1, -1)

        # Zone mask
        zone_mask = torch.ones(B, self.num_regions, S, device=x.device, dtype=torch.bool)
        for b in range(B):
            n = int(mask[b].sum().item()) if mask is not None else S
            zone_size = n / self.num_zones
            for z in range(self.num_zones):
                s, e = int(round(z * zone_size)), min(max(int(round((z + 1) * zone_size)), int(round(z * zone_size)) + 1), n)
                q_s, q_e = z * self.tokens_per_zone, (z + 1) * self.tokens_per_zone
                zone_mask[b, q_s:q_e, s:e] = False
        zone_mask = zone_mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1).reshape(B * self.num_heads, self.num_regions, S)

        kpm = (mask == 0) if mask is not None else None
        region_hidden, attn_w = self.attention(queries, x_proj, x_proj, attn_mask=zone_mask, key_padding_mask=kpm, need_weights=True)
        return self.output_proj(region_hidden), attn_w


class VisualEncoder(nn.Module):
    """Compresses variable-length slice embeddings into fixed visual tokens in LLM space."""

    def __init__(self, slice_dim=1024, hidden_dim=1024, llm_dim=3072,
                 num_regions=32, num_zones=32, tokens_per_zone=1,
                 num_heads=16, dropout=0.1):
        super().__init__()
        self.region_query = ZonedCrossAttention(slice_dim, hidden_dim, num_zones, tokens_per_zone, num_heads, dropout)
        self.global_self_attn = nn.TransformerEncoderLayer(
            d_model=slice_dim, nhead=num_heads, batch_first=True,
            dropout=dropout, dim_feedforward=slice_dim * 4, activation="gelu",
        )
        self.jepa_predictor = nn.Sequential(
            nn.Dropout(0.0),  # disabled at inference
            nn.Linear(slice_dim, llm_dim),
            nn.LayerNorm(llm_dim),
            nn.Dropout(0.0),
        )
        self.norm_calibrator_scale = 1.0  # set from checkpoint buffer
        self.register_buffer("_norm_scale", torch.tensor(1.0))

    def forward(self, slices, mask=None):
        """
        Args:
            slices: (B, num_slices, 1024) pre-computed LeJEPA embeddings
            mask:   (B, num_slices) binary mask, 1=valid 0=pad
        Returns:
            (B, 32, 3072) visual tokens in LLM hidden space
        """
        regions, _ = self.region_query(slices, mask)
        regions = self.global_self_attn(regions)
        predicted = self.jepa_predictor(regions)
        return predicted * self._norm_scale


class GatedCrossAttentionLayer(nn.Module):
    """Flamingo-style cross-attention: text hidden states attend to visual tokens."""

    def __init__(self, hidden_dim=3072, num_heads=16):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.o_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.q_norm = nn.LayerNorm(hidden_dim)
        self.gate = nn.Parameter(torch.tensor(0.0), requires_grad=False)

    def forward(self, text_hidden, visual_tokens):
        B, S, _ = text_hidden.shape
        dt = self.q_proj.weight.dtype
        text_hidden = text_hidden.to(dt)
        visual_tokens = visual_tokens.to(dt)
        q = self.q_proj(self.q_norm(text_hidden)).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(visual_tokens).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(visual_tokens).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
        return self.o_proj(out.transpose(1, 2).contiguous().view(B, S, self.hidden_dim))


# ---------------------------------------------------------------------------
# Main inference model
# ---------------------------------------------------------------------------

class KerVLJEPA(nn.Module):
    """
    Ker-VLJEPA-3B inference model.
    Takes pre-computed LeJEPA slice embeddings (1024-d per slice) and generates
    free-text radiology reports using Llama 3.2 3B + LoRA + cross-attention bridge.
    """

    INJECTION_LAYERS = [7, 14, 21]
    NUM_REGIONS = 32
    VISUAL_TOKEN = "<|visual_region|>"

    def __init__(self, llm_path: str, weights_dir: str = "weights", device: str = "cuda"):
        super().__init__()
        weights_dir = Path(weights_dir)
        self.device = torch.device(device)

        # --- 1. Load tokenizer ---
        tok_dir = weights_dir / "tokenizer"
        self.tokenizer = AutoTokenizer.from_pretrained(
            llm_path, trust_remote_code=True, padding_side="left",
        )
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.add_tokens([self.VISUAL_TOKEN], special_tokens=True)
        self.vis_token_id = self.tokenizer.convert_tokens_to_ids(self.VISUAL_TOKEN)

        # Load custom chat template
        tmpl_path = tok_dir / "chat_template.jinja"
        if tmpl_path.exists():
            self.tokenizer.chat_template = tmpl_path.read_text()

        # --- 2. Load Llama 3.2 3B + LoRA ---
        self.llm = AutoModelForCausalLM.from_pretrained(
            llm_path, torch_dtype=torch.bfloat16,
            device_map=device, trust_remote_code=True,
            attn_implementation="flash_attention_2",
        )
        self.llm.resize_token_embeddings(len(self.tokenizer))

        lora_dir = weights_dir / "lora_adapters"
        self.llm = PeftModel.from_pretrained(self.llm, str(lora_dir), is_trainable=False)
        self.llm.eval()

        llm_dim = self.llm.config.hidden_size  # 3072

        # --- 3. Visual encoder ---
        self.visual_encoder = VisualEncoder(
            slice_dim=1024, hidden_dim=1024, llm_dim=llm_dim,
            num_regions=self.NUM_REGIONS, num_zones=32, tokens_per_zone=1,
            num_heads=16, dropout=0.0,
        )
        ve_state = load_file(str(weights_dir / "visual_encoder.safetensors"))
        # Map checkpoint keys (which include output_ln, norm_calibrator, classifiers)
        # into our simplified VisualEncoder
        mapped = {}
        for k, v in ve_state.items():
            if k == "norm_calibrator.scale":
                self.visual_encoder._norm_scale = v.clone()
                continue
            if k.startswith("norm_calibrator.") or k.startswith("region_classifier.") or \
               k.startswith("slice_organ_classifier.") or k == "_last_attention_weights" or \
               k.startswith("output_ln."):
                continue
            mapped[k] = v
        self.visual_encoder.load_state_dict(mapped, strict=False)
        self.visual_encoder = self.visual_encoder.to(self.device).to(torch.bfloat16)
        self.visual_encoder.eval()

        # --- 4. Bridge components ---
        bridge = load_file(str(weights_dir / "bridge_components.safetensors"))

        # Text embedding norm (for grafting normalization)
        self.register_buffer("text_embed_norm", bridge["text_embed_norm"].to(self.device))

        # LayerNorm for grafting
        self.layernorm = nn.LayerNorm(llm_dim).to(self.device).to(torch.bfloat16)
        self.layernorm.weight.data.copy_(bridge["layernorm.weight"])
        self.layernorm.bias.data.copy_(bridge["layernorm.bias"])

        # Cross-attention adapters + layer projectors
        self.cross_attn_adapters = nn.ModuleDict()
        self.layer_projectors = nn.ModuleDict()
        for layer_idx in self.INJECTION_LAYERS:
            li = str(layer_idx)
            adapter = GatedCrossAttentionLayer(llm_dim, 16)
            adapter.q_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.q_proj.weight"])
            adapter.k_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.k_proj.weight"])
            adapter.v_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.v_proj.weight"])
            adapter.o_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.o_proj.weight"])
            adapter.q_norm.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.q_norm.weight"])
            adapter.q_norm.bias.data.copy_(bridge[f"cross_attn_adapters.{li}.q_norm.bias"])
            self.cross_attn_adapters[li] = adapter.to(self.device).to(torch.bfloat16)

            proj = nn.Linear(llm_dim, llm_dim, bias=False)
            proj.weight.data.copy_(bridge[f"layer_projectors.{li}.weight"])
            self.layer_projectors[li] = proj.to(self.device).to(torch.bfloat16)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _get_llm_layers(self):
        model = self.llm
        for _ in range(10):
            if hasattr(model, "base_model"):
                model = model.base_model
            elif hasattr(model, "model"):
                model = model.model
            else:
                break
        return model.layers

    def _get_embed_layer(self):
        model = self.llm
        for _ in range(10):
            if hasattr(model, "get_input_embeddings"):
                emb = model.get_input_embeddings()
                if emb is not None:
                    return emb
            if hasattr(model, "base_model"):
                model = model.base_model
            elif hasattr(model, "model"):
                model = model.model
            else:
                break
        raise RuntimeError("Could not find embedding layer")

    def _normalize_visual(self, visual_embeds):
        norms = visual_embeds.norm(dim=-1, keepdim=True).clamp(min=1e-8)
        return visual_embeds / norms * self.text_embed_norm

    @staticmethod
    def _clean_text(text: str) -> str:
        text = text.strip()
        if not text:
            return text
        low = text.lower()
        first = low.find("findings:")
        if first >= 0:
            second = low.find("findings:", first + 10)
            if second > 0:
                text = text[:second].strip()
        for marker in ["\n\n\n", "User", "user", "assistant", "system"]:
            idx = text.find(marker)
            if idx > 20:
                text = text[:idx].strip()
        return text

    # ------------------------------------------------------------------
    # Generation
    # ------------------------------------------------------------------

    @torch.no_grad()
    def generate(
        self,
        slice_embeddings: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        max_new_tokens: int = 384,
        temperature: float = 0.6,
        top_p: float = 0.9,
        repetition_penalty: float = 1.1,
        no_repeat_ngram_size: int = 4,
    ) -> str:
        """
        Generate a radiology report from pre-computed LeJEPA slice embeddings.
        Args:
            slice_embeddings: (1, num_slices, 1024) — stacked 1024-d embeddings,
                              one per CT slice. Padding is allowed.
            mask: (1, num_slices) — binary mask where 1=real slice, 0=padding.
                  If None, all slices are treated as valid.
            max_new_tokens: maximum tokens to generate (default 384).
            temperature: sampling temperature (default 0.6).
            top_p: nucleus sampling threshold (default 0.9).
            repetition_penalty: penalize repeated tokens (default 1.1).
            no_repeat_ngram_size: prevent repeating n-grams (default 4).
        Returns:
            Generated report text (str).
        """
        assert slice_embeddings.ndim == 3 and slice_embeddings.shape[0] == 1

        # 1. Visual forward
        slices = slice_embeddings.to(self.device, dtype=torch.bfloat16)
        if mask is not None:
            mask = mask.to(self.device, dtype=torch.bfloat16)
        visual_tokens = self.visual_encoder(slices, mask)
        visual_tokens = self._normalize_visual(visual_tokens.to(torch.bfloat16))

        # 2. Build prompt
        placeholders = self.VISUAL_TOKEN * self.NUM_REGIONS
        messages = [
            {"role": "system", "content": "You are a radiology reporting assistant. Describe thoracic findings based on the provided CT scan visual features. Report only what you observe."},
            {"role": "user", "content": f"Based on the visual features from this CT scan, describe the thoracic findings. {placeholders}"},
        ]
        tokenized = self.tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
        )
        if hasattr(tokenized, "input_ids"):
            input_ids = tokenized["input_ids"].to(self.device)
            attention_mask = tokenized["attention_mask"].to(self.device)
        else:
            input_ids = tokenized.to(self.device)
            attention_mask = torch.ones_like(input_ids)

        # 3. Graft visual tokens into embedding sequence
        embed_layer = self._get_embed_layer()
        inputs_embeds = embed_layer(input_ids).clone()
        vis_mask = (input_ids == self.vis_token_id)
        vis_positions = vis_mask[0].nonzero(as_tuple=True)[0]
        assert len(vis_positions) == self.NUM_REGIONS, \
            f"Expected {self.NUM_REGIONS} visual tokens, found {len(vis_positions)}"
        for idx, pos in enumerate(vis_positions):
            inputs_embeds[0, pos] = visual_tokens[0, idx].to(inputs_embeds.dtype)

        # 4. Register cross-attention hooks
        hooks = []
        llm_layers = self._get_llm_layers()
        for layer_idx in self.INJECTION_LAYERS:
            li = str(layer_idx)
            proj = self.layer_projectors[li]
            adapter = self.cross_attn_adapters[li]

            def make_hook(p, a, v_tokens, v_mask, seq_len):
                def hook_fn(module, args, output):
                    hidden = output[0] if isinstance(output, tuple) else output
                    # Additive injection at visual positions (prefill only)
                    if hidden.shape[1] == seq_len:
                        projected = p(v_tokens.to(hidden.dtype))
                        vis_pos = v_mask[0].nonzero(as_tuple=True)[0]
                        hidden[0, vis_pos] = hidden[0, vis_pos] + projected[0, :len(vis_pos)]
                    # Cross-attention (every token)
                    xattn_out = a(hidden, v_tokens.to(hidden.dtype))
                    modified = hidden + xattn_out.to(hidden.dtype)
                    return (modified,) + output[1:] if isinstance(output, tuple) else modified
                return hook_fn

            h = llm_layers[layer_idx].register_forward_hook(
                make_hook(proj, adapter, visual_tokens, vis_mask, input_ids.shape[1])
            )
            hooks.append(h)

        # 5. Generate
        eot_id = self.tokenizer.convert_tokens_to_ids("<|eot_id|>")
        start_header = self.tokenizer.convert_tokens_to_ids("<|start_header_id|>")
        stop_ids = [self.tokenizer.eos_token_id]
        if eot_id is not None and eot_id != self.tokenizer.eos_token_id:
            stop_ids.append(eot_id)
        if start_header is not None and start_header not in stop_ids:
            stop_ids.append(start_header)

        try:
            generated_ids = self.llm.generate(
                inputs_embeds=inputs_embeds,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                do_sample=True,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=stop_ids,
            )
        finally:
            for h in hooks:
                h.remove()

        text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        return self._clean_text(text)

    @torch.no_grad()
    def classify(
        self,
        slice_embeddings: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> Dict[str, float]:
        """
        Run the auxiliary 18-class abnormality classifier on slice embeddings.
        Returns a dict mapping each CT-RATE condition name to its sigmoid probability.
        """
        CLASS_NAMES = [
            "Medical material", "Arterial wall calcification",
            "Cardiomegaly", "Pericardial effusion",
            "Coronary artery wall calcification", "Hiatal hernia",
            "Lymphadenopathy", "Emphysema",
            "Atelectasis", "Lung nodule",
            "Lung opacity", "Pulmonary fibrotic sequela",
            "Pleural effusion", "Mosaic attenuation pattern",
            "Peribronchial thickening", "Consolidation",
            "Bronchiectasis", "Interlobular septal thickening",
        ]
        slices = slice_embeddings.to(self.device, dtype=torch.bfloat16)
        if mask is not None:
            mask = mask.to(self.device, dtype=torch.bfloat16)
        visual_tokens = self.visual_encoder(slices, mask)  # (1, 32, 3072)
        pooled = visual_tokens.mean(dim=1)                 # (1, 3072)
        # case_classifier is loaded from bridge but we need to add it
        logits = self._case_classifier(pooled)              # (1, 18)
        probs = torch.sigmoid(logits).squeeze(0).cpu().tolist()
        return {name: round(p, 4) for name, p in zip(CLASS_NAMES, probs)}


def load_model(llm_path: str, weights_dir: str = "weights", device: str = "cuda") -> KerVLJEPA:
    """
    Load the Ker-VLJEPA-3B model for inference.
    Args:
        llm_path: Path to local Llama 3.2 3B model directory.
        weights_dir: Path to the weights/ directory from this package.
        device: CUDA device string (default "cuda").
    Returns:
        Ready-to-use KerVLJEPA model instance.
    """
    model = KerVLJEPA(llm_path=llm_path, weights_dir=weights_dir, device=device)
    model.eval()
    return model

	"""
	Ker-VLJEPA-3B — Inference-only model.

	Loads the Llama 3.2 3B backbone, applies LoRA adapters, visual encoder,
	and cross-attention bridge components for CT report generation.

	Requires:
	- A local copy of meta-llama/Llama-3.2-3B (user-provided)
	- Weight files in weights/ (shipped with this package)
	"""
	import math
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from pathlib import Path
	from typing import Optional, List, Tuple, Dict
	from transformers import AutoModelForCausalLM, AutoTokenizer
	from peft import PeftModel
	from safetensors.torch import load_file


	# ---------------------------------------------------------------------------
	# Visual encoder components (inference-only, no training heads)
	# ---------------------------------------------------------------------------

	def sinusoidal_positional_encoding(n: int, d: int) -> torch.Tensor:
	pe = torch.zeros(n, d)
	pos = torch.arange(n, dtype=torch.float).unsqueeze(1)
	div = torch.exp(torch.arange(0, d, 2).float() * (-math.log(10000.0) / d))
	pe[:, 0::2] = torch.sin(pos * div)
	pe[:, 1::2] = torch.cos(pos * div)
	return pe.unsqueeze(0)


	class ZonedCrossAttention(nn.Module):
	"""Z-Zoned cross-attention: each zone's queries attend only to their spatial slice range."""

	def __init__(self, slice_dim=1024, hidden_dim=1024, num_zones=32,
	tokens_per_zone=1, num_heads=16, dropout=0.1):
	super().__init__()
	self.num_zones = num_zones
	self.tokens_per_zone = tokens_per_zone
	self.num_regions = num_zones * tokens_per_zone
	self.num_heads = num_heads
	self.hidden_dim = hidden_dim

	self.zone_queries = nn.Parameter(torch.zeros(num_zones, tokens_per_zone, hidden_dim))
	self.zone_embed = nn.Embedding(num_zones, hidden_dim)
	self.slice_proj = nn.Linear(slice_dim, hidden_dim)
	self.attention = nn.MultiheadAttention(hidden_dim, num_heads, batch_first=True, dropout=dropout)
	self.output_proj = nn.Linear(hidden_dim, slice_dim)
	self.fallback_pos_embed = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02)

	# Physical Z-positional encoding buffers
	div_term = torch.exp(torch.arange(0, hidden_dim, 2).float() * (-math.log(10000.0) / hidden_dim))
	self.register_buffer("div_term", div_term)

	def _add_z_pos(self, x_proj, mask):
	"""Sinusoidal positional encoding based on slice index."""
	B, S, D = x_proj.shape
	device, dtype = x_proj.device, x_proj.dtype
	z_positions = torch.zeros(B, S, device=device)
	for b in range(B):
	n = int(mask[b].sum().item()) if mask is not None else S
	pos = torch.arange(n, device=device, dtype=torch.float) * 2.5 / 600.0 * 100.0
	z_positions[b, :n] = pos
	pos = z_positions.unsqueeze(-1)
	pe = torch.zeros(B, S, D, device=device, dtype=dtype)
	pe[:, :, 0::2] = torch.sin(pos * self.div_term)
	pe[:, :, 1::2] = torch.cos(pos * self.div_term)
	if mask is not None:
	pe = pe * mask.unsqueeze(-1)
	return x_proj + pe

	def forward(self, x, mask=None, metadata=None):
	B, S, _ = x.shape
	x_proj = self.slice_proj(x)
	if mask is not None:
	x_proj = self._add_z_pos(x_proj, mask)
	else:
	x_proj = x_proj + self.fallback_pos_embed

	queries = self.zone_queries.reshape(self.num_regions, self.hidden_dim)
	zone_ids = torch.arange(self.num_zones, device=x.device)
	zone_emb = self.zone_embed(zone_ids).unsqueeze(1).expand(-1, self.tokens_per_zone, -1)
	queries = (queries + zone_emb.reshape(self.num_regions, self.hidden_dim)).unsqueeze(0).expand(B, -1, -1)

	# Zone mask
	zone_mask = torch.ones(B, self.num_regions, S, device=x.device, dtype=torch.bool)
	for b in range(B):
	n = int(mask[b].sum().item()) if mask is not None else S
	zone_size = n / self.num_zones
	for z in range(self.num_zones):
	s, e = int(round(z * zone_size)), min(max(int(round((z + 1) * zone_size)), int(round(z * zone_size)) + 1), n)
	q_s, q_e = z * self.tokens_per_zone, (z + 1) * self.tokens_per_zone
	zone_mask[b, q_s:q_e, s:e] = False
	zone_mask = zone_mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1).reshape(B * self.num_heads, self.num_regions, S)

	kpm = (mask == 0) if mask is not None else None
	region_hidden, attn_w = self.attention(queries, x_proj, x_proj, attn_mask=zone_mask, key_padding_mask=kpm, need_weights=True)
	return self.output_proj(region_hidden), attn_w


	class VisualEncoder(nn.Module):
	"""Compresses variable-length slice embeddings into fixed visual tokens in LLM space."""

	def __init__(self, slice_dim=1024, hidden_dim=1024, llm_dim=3072,
	num_regions=32, num_zones=32, tokens_per_zone=1,
	num_heads=16, dropout=0.1):
	super().__init__()
	self.region_query = ZonedCrossAttention(slice_dim, hidden_dim, num_zones, tokens_per_zone, num_heads, dropout)
	self.global_self_attn = nn.TransformerEncoderLayer(
	d_model=slice_dim, nhead=num_heads, batch_first=True,
	dropout=dropout, dim_feedforward=slice_dim * 4, activation="gelu",
	)
	self.jepa_predictor = nn.Sequential(
	nn.Dropout(0.0), # disabled at inference
	nn.Linear(slice_dim, llm_dim),
	nn.LayerNorm(llm_dim),
	nn.Dropout(0.0),
	)
	self.norm_calibrator_scale = 1.0 # set from checkpoint buffer
	self.register_buffer("_norm_scale", torch.tensor(1.0))

	def forward(self, slices, mask=None):
	"""
	Args:
	slices: (B, num_slices, 1024) pre-computed LeJEPA embeddings
	mask: (B, num_slices) binary mask, 1=valid 0=pad
	Returns:
	(B, 32, 3072) visual tokens in LLM hidden space
	"""
	regions, _ = self.region_query(slices, mask)
	regions = self.global_self_attn(regions)
	predicted = self.jepa_predictor(regions)
	return predicted * self._norm_scale


	class GatedCrossAttentionLayer(nn.Module):
	"""Flamingo-style cross-attention: text hidden states attend to visual tokens."""

	def __init__(self, hidden_dim=3072, num_heads=16):
	super().__init__()
	self.hidden_dim = hidden_dim
	self.num_heads = num_heads
	self.head_dim = hidden_dim // num_heads
	self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
	self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
	self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
	self.o_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
	self.q_norm = nn.LayerNorm(hidden_dim)
	self.gate = nn.Parameter(torch.tensor(0.0), requires_grad=False)

	def forward(self, text_hidden, visual_tokens):
	B, S, _ = text_hidden.shape
	dt = self.q_proj.weight.dtype
	text_hidden = text_hidden.to(dt)
	visual_tokens = visual_tokens.to(dt)
	q = self.q_proj(self.q_norm(text_hidden)).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
	k = self.k_proj(visual_tokens).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
	v = self.v_proj(visual_tokens).view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
	out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
	return self.o_proj(out.transpose(1, 2).contiguous().view(B, S, self.hidden_dim))


	# ---------------------------------------------------------------------------
	# Main inference model
	# ---------------------------------------------------------------------------

	class KerVLJEPA(nn.Module):
	"""
	Ker-VLJEPA-3B inference model.

	Takes pre-computed LeJEPA slice embeddings (1024-d per slice) and generates
	free-text radiology reports using Llama 3.2 3B + LoRA + cross-attention bridge.
	"""

	INJECTION_LAYERS = [7, 14, 21]
	NUM_REGIONS = 32
	VISUAL_TOKEN = "<\|visual_region\|>"

	def __init__(self, llm_path: str, weights_dir: str = "weights", device: str = "cuda"):
	super().__init__()
	weights_dir = Path(weights_dir)
	self.device = torch.device(device)

	# --- 1. Load tokenizer ---
	tok_dir = weights_dir / "tokenizer"
	self.tokenizer = AutoTokenizer.from_pretrained(
	llm_path, trust_remote_code=True, padding_side="left",
	)
	if self.tokenizer.pad_token is None:
	self.tokenizer.pad_token = self.tokenizer.eos_token
	self.tokenizer.add_tokens([self.VISUAL_TOKEN], special_tokens=True)
	self.vis_token_id = self.tokenizer.convert_tokens_to_ids(self.VISUAL_TOKEN)

	# Load custom chat template
	tmpl_path = tok_dir / "chat_template.jinja"
	if tmpl_path.exists():
	self.tokenizer.chat_template = tmpl_path.read_text()

	# --- 2. Load Llama 3.2 3B + LoRA ---
	self.llm = AutoModelForCausalLM.from_pretrained(
	llm_path, torch_dtype=torch.bfloat16,
	device_map=device, trust_remote_code=True,
	attn_implementation="flash_attention_2",
	)
	self.llm.resize_token_embeddings(len(self.tokenizer))

	lora_dir = weights_dir / "lora_adapters"
	self.llm = PeftModel.from_pretrained(self.llm, str(lora_dir), is_trainable=False)
	self.llm.eval()

	llm_dim = self.llm.config.hidden_size # 3072

	# --- 3. Visual encoder ---
	self.visual_encoder = VisualEncoder(
	slice_dim=1024, hidden_dim=1024, llm_dim=llm_dim,
	num_regions=self.NUM_REGIONS, num_zones=32, tokens_per_zone=1,
	num_heads=16, dropout=0.0,
	)
	ve_state = load_file(str(weights_dir / "visual_encoder.safetensors"))
	# Map checkpoint keys (which include output_ln, norm_calibrator, classifiers)
	# into our simplified VisualEncoder
	mapped = {}
	for k, v in ve_state.items():
	if k == "norm_calibrator.scale":
	self.visual_encoder._norm_scale = v.clone()
	continue
	if k.startswith("norm_calibrator.") or k.startswith("region_classifier.") or \
	k.startswith("slice_organ_classifier.") or k == "_last_attention_weights" or \
	k.startswith("output_ln."):
	continue
	mapped[k] = v
	self.visual_encoder.load_state_dict(mapped, strict=False)
	self.visual_encoder = self.visual_encoder.to(self.device).to(torch.bfloat16)
	self.visual_encoder.eval()

	# --- 4. Bridge components ---
	bridge = load_file(str(weights_dir / "bridge_components.safetensors"))

	# Text embedding norm (for grafting normalization)
	self.register_buffer("text_embed_norm", bridge["text_embed_norm"].to(self.device))

	# LayerNorm for grafting
	self.layernorm = nn.LayerNorm(llm_dim).to(self.device).to(torch.bfloat16)
	self.layernorm.weight.data.copy_(bridge["layernorm.weight"])
	self.layernorm.bias.data.copy_(bridge["layernorm.bias"])

	# Cross-attention adapters + layer projectors
	self.cross_attn_adapters = nn.ModuleDict()
	self.layer_projectors = nn.ModuleDict()
	for layer_idx in self.INJECTION_LAYERS:
	li = str(layer_idx)
	adapter = GatedCrossAttentionLayer(llm_dim, 16)
	adapter.q_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.q_proj.weight"])
	adapter.k_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.k_proj.weight"])
	adapter.v_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.v_proj.weight"])
	adapter.o_proj.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.o_proj.weight"])
	adapter.q_norm.weight.data.copy_(bridge[f"cross_attn_adapters.{li}.q_norm.weight"])
	adapter.q_norm.bias.data.copy_(bridge[f"cross_attn_adapters.{li}.q_norm.bias"])
	self.cross_attn_adapters[li] = adapter.to(self.device).to(torch.bfloat16)

	proj = nn.Linear(llm_dim, llm_dim, bias=False)
	proj.weight.data.copy_(bridge[f"layer_projectors.{li}.weight"])
	self.layer_projectors[li] = proj.to(self.device).to(torch.bfloat16)

	# ------------------------------------------------------------------
	# Helpers
	# ------------------------------------------------------------------

	def _get_llm_layers(self):
	model = self.llm
	for _ in range(10):
	if hasattr(model, "base_model"):
	model = model.base_model
	elif hasattr(model, "model"):
	model = model.model
	else:
	break
	return model.layers

	def _get_embed_layer(self):
	model = self.llm
	for _ in range(10):
	if hasattr(model, "get_input_embeddings"):
	emb = model.get_input_embeddings()
	if emb is not None:
	return emb
	if hasattr(model, "base_model"):
	model = model.base_model
	elif hasattr(model, "model"):
	model = model.model
	else:
	break
	raise RuntimeError("Could not find embedding layer")

	def _normalize_visual(self, visual_embeds):
	norms = visual_embeds.norm(dim=-1, keepdim=True).clamp(min=1e-8)
	return visual_embeds / norms * self.text_embed_norm

	@staticmethod
	def _clean_text(text: str) -> str:
	text = text.strip()
	if not text:
	return text
	low = text.lower()
	first = low.find("findings:")
	if first >= 0:
	second = low.find("findings:", first + 10)
	if second > 0:
	text = text[:second].strip()
	for marker in ["\n\n\n", "User", "user", "assistant", "system"]:
	idx = text.find(marker)
	if idx > 20:
	text = text[:idx].strip()
	return text

	# ------------------------------------------------------------------
	# Generation
	# ------------------------------------------------------------------

	@torch.no_grad()
	def generate(
	self,
	slice_embeddings: torch.Tensor,
	mask: Optional[torch.Tensor] = None,
	max_new_tokens: int = 384,
	temperature: float = 0.6,
	top_p: float = 0.9,
	repetition_penalty: float = 1.1,
	no_repeat_ngram_size: int = 4,
	) -> str:
	"""
	Generate a radiology report from pre-computed LeJEPA slice embeddings.

	Args:
	slice_embeddings: (1, num_slices, 1024) — stacked 1024-d embeddings,
	one per CT slice. Padding is allowed.
	mask: (1, num_slices) — binary mask where 1=real slice, 0=padding.
	If None, all slices are treated as valid.
	max_new_tokens: maximum tokens to generate (default 384).
	temperature: sampling temperature (default 0.6).
	top_p: nucleus sampling threshold (default 0.9).
	repetition_penalty: penalize repeated tokens (default 1.1).
	no_repeat_ngram_size: prevent repeating n-grams (default 4).

	Returns:
	Generated report text (str).
	"""
	assert slice_embeddings.ndim == 3 and slice_embeddings.shape[0] == 1

	# 1. Visual forward
	slices = slice_embeddings.to(self.device, dtype=torch.bfloat16)
	if mask is not None:
	mask = mask.to(self.device, dtype=torch.bfloat16)
	visual_tokens = self.visual_encoder(slices, mask)
	visual_tokens = self._normalize_visual(visual_tokens.to(torch.bfloat16))

	# 2. Build prompt
	placeholders = self.VISUAL_TOKEN * self.NUM_REGIONS
	messages = [
	{"role": "system", "content": "You are a radiology reporting assistant. Describe thoracic findings based on the provided CT scan visual features. Report only what you observe."},
	{"role": "user", "content": f"Based on the visual features from this CT scan, describe the thoracic findings. {placeholders}"},
	]
	tokenized = self.tokenizer.apply_chat_template(
	messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
	)
	if hasattr(tokenized, "input_ids"):
	input_ids = tokenized["input_ids"].to(self.device)
	attention_mask = tokenized["attention_mask"].to(self.device)
	else:
	input_ids = tokenized.to(self.device)
	attention_mask = torch.ones_like(input_ids)

	# 3. Graft visual tokens into embedding sequence
	embed_layer = self._get_embed_layer()
	inputs_embeds = embed_layer(input_ids).clone()
	vis_mask = (input_ids == self.vis_token_id)
	vis_positions = vis_mask[0].nonzero(as_tuple=True)[0]
	assert len(vis_positions) == self.NUM_REGIONS, \
	f"Expected {self.NUM_REGIONS} visual tokens, found {len(vis_positions)}"
	for idx, pos in enumerate(vis_positions):
	inputs_embeds[0, pos] = visual_tokens[0, idx].to(inputs_embeds.dtype)

	# 4. Register cross-attention hooks
	hooks = []
	llm_layers = self._get_llm_layers()
	for layer_idx in self.INJECTION_LAYERS:
	li = str(layer_idx)
	proj = self.layer_projectors[li]
	adapter = self.cross_attn_adapters[li]

	def make_hook(p, a, v_tokens, v_mask, seq_len):
	def hook_fn(module, args, output):
	hidden = output[0] if isinstance(output, tuple) else output
	# Additive injection at visual positions (prefill only)
	if hidden.shape[1] == seq_len:
	projected = p(v_tokens.to(hidden.dtype))
	vis_pos = v_mask[0].nonzero(as_tuple=True)[0]
	hidden[0, vis_pos] = hidden[0, vis_pos] + projected[0, :len(vis_pos)]
	# Cross-attention (every token)
	xattn_out = a(hidden, v_tokens.to(hidden.dtype))
	modified = hidden + xattn_out.to(hidden.dtype)
	return (modified,) + output[1:] if isinstance(output, tuple) else modified
	return hook_fn

	h = llm_layers[layer_idx].register_forward_hook(
	make_hook(proj, adapter, visual_tokens, vis_mask, input_ids.shape[1])
	)
	hooks.append(h)

	# 5. Generate
	eot_id = self.tokenizer.convert_tokens_to_ids("<\|eot_id\|>")
	start_header = self.tokenizer.convert_tokens_to_ids("<\|start_header_id\|>")
	stop_ids = [self.tokenizer.eos_token_id]
	if eot_id is not None and eot_id != self.tokenizer.eos_token_id:
	stop_ids.append(eot_id)
	if start_header is not None and start_header not in stop_ids:
	stop_ids.append(start_header)

	try:
	generated_ids = self.llm.generate(
	inputs_embeds=inputs_embeds,
	attention_mask=attention_mask,
	max_new_tokens=max_new_tokens,
	temperature=temperature,
	top_p=top_p,
	repetition_penalty=repetition_penalty,
	no_repeat_ngram_size=no_repeat_ngram_size,
	do_sample=True,
	pad_token_id=self.tokenizer.pad_token_id,
	eos_token_id=stop_ids,
	)
	finally:
	for h in hooks:
	h.remove()

	text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
	return self._clean_text(text)

	@torch.no_grad()
	def classify(
	self,
	slice_embeddings: torch.Tensor,
	mask: Optional[torch.Tensor] = None,
	) -> Dict[str, float]:
	"""
	Run the auxiliary 18-class abnormality classifier on slice embeddings.

	Returns a dict mapping each CT-RATE condition name to its sigmoid probability.
	"""
	CLASS_NAMES = [
	"Medical material", "Arterial wall calcification",
	"Cardiomegaly", "Pericardial effusion",
	"Coronary artery wall calcification", "Hiatal hernia",
	"Lymphadenopathy", "Emphysema",
	"Atelectasis", "Lung nodule",
	"Lung opacity", "Pulmonary fibrotic sequela",
	"Pleural effusion", "Mosaic attenuation pattern",
	"Peribronchial thickening", "Consolidation",
	"Bronchiectasis", "Interlobular septal thickening",
	]
	slices = slice_embeddings.to(self.device, dtype=torch.bfloat16)
	if mask is not None:
	mask = mask.to(self.device, dtype=torch.bfloat16)
	visual_tokens = self.visual_encoder(slices, mask) # (1, 32, 3072)
	pooled = visual_tokens.mean(dim=1) # (1, 3072)
	# case_classifier is loaded from bridge but we need to add it
	logits = self._case_classifier(pooled) # (1, 18)
	probs = torch.sigmoid(logits).squeeze(0).cpu().tolist()
	return {name: round(p, 4) for name, p in zip(CLASS_NAMES, probs)}


	def load_model(llm_path: str, weights_dir: str = "weights", device: str = "cuda") -> KerVLJEPA:
	"""
	Load the Ker-VLJEPA-3B model for inference.

	Args:
	llm_path: Path to local Llama 3.2 3B model directory.
	weights_dir: Path to the weights/ directory from this package.
	device: CUDA device string (default "cuda").

	Returns:
	Ready-to-use KerVLJEPA model instance.
	"""
	model = KerVLJEPA(llm_path=llm_path, weights_dir=weights_dir, device=device)
	model.eval()
	return model