source_gemma.py

"""Strict persistent-visual CVRR wrapper for Gemma 3 and Gemma 4.

The implementation intentionally wraps released Transformers models instead
of copying their decoder source.  Native layer-call arguments are captured
during frozen prefix passes and replayed for the shared recurrent cell and
upper text-only continuation.  This preserves each release's masks, rotary
geometry, and attention type while enforcing a visibly auditable interface:
the upper decoder receives only non-visual recurrent rows.
"""

from __future__ import annotations

import contextlib
import io
import json
import math
import pathlib
from dataclasses import dataclass
from typing import Any

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

from .source_helpers import _layer_hidden


class LoRALinear(nn.Module):
    def __init__(
        self,
        base: nn.Linear,
        *,
        rank: int,
        alpha: float,
        dropout: float,
    ):
        super().__init__()
        if rank <= 0:
            raise ValueError("LoRA rank must be positive")
        self.base = base
        for parameter in self.base.parameters():
            parameter.requires_grad_(False)
        self.rank = int(rank)
        self.scale = float(alpha) / float(rank)
        self.dropout = nn.Dropout(float(dropout))
        self.lora_A = nn.Parameter(torch.empty(rank, base.in_features))
        self.lora_B = nn.Parameter(torch.zeros(base.out_features, rank))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        self.enabled = True

    def forward(self, inputs):
        result = self.base(inputs)
        if not self.enabled:
            return result
        update = F.linear(self.dropout(inputs).float(), self.lora_A.float())
        update = F.linear(update, self.lora_B.float())
        return result + (update * self.scale).to(result.dtype)


def _module_parent(root: nn.Module, path: str):
    parts = path.split(".")
    parent = root
    for part in parts[:-1]:
        parent = getattr(parent, part)
    return parent, parts[-1]


def inject_cell_lora(
    cell: nn.Module,
    *,
    rank: int,
    alpha: float,
    dropout: float,
    suffixes: set[str] | None = None,
) -> dict[str, LoRALinear]:
    if suffixes is None:
        suffixes = {
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        }
    selected = {
        name: module
        for name, module in cell.named_modules()
        if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in suffixes
    }
    if not selected:
        raise RuntimeError("no attention/MLP projections found in recurrent cell")
    wrappers = {}
    # Materialize the list before replacing children during traversal.
    for name, module in selected.items():
        parent, child = _module_parent(cell, name)
        wrapper = LoRALinear(
            module, rank=rank, alpha=alpha, dropout=dropout
        )
        setattr(parent, child, wrapper)
        wrappers[name] = wrapper
    return wrappers


@dataclass
class LayerCall:
    positional_tail: tuple[Any, ...]
    keywords: dict[str, Any]


@dataclass
class GemmaCVRRTrace:
    """Frozen native states needed for one strict recurrent rollout.

    ``scaffold`` is multimodal and is consumed only by the shared recurrent
    cell.  ``text_calls`` comes from an image-free sequence and is the sole
    context replayed by the upper decoder.  Keeping those two objects separate
    makes the no-bypass contract directly inspectable.
    """

    scaffold: torch.Tensor
    text_rows: torch.Tensor
    visual_rows: torch.Tensor
    question_valid: torch.Tensor
    base_anchor: torch.Tensor
    r1: torch.Tensor
    mm_cell_call: LayerCall
    text_calls: dict[int, LayerCall]

    @property
    def question_lengths(self) -> torch.Tensor:
        return self.question_valid.sum(dim=1)

    @property
    def visual_lengths(self) -> torch.Tensor:
        return self.visual_rows.sum(dim=1)


class _StopAfterCell(RuntimeError):
    pass


def _capture_call(storage: dict[int, LayerCall], index: int):
    def hook(_module, args, kwargs):
        storage[index] = LayerCall(tuple(args[1:]), dict(kwargs))

    return hook


def _gather_rows(full, mask):
    batch, _, width = full.shape
    lengths = mask.sum(dim=1).tolist()
    result = full.new_zeros((batch, max(lengths), width))
    for row, length in enumerate(lengths):
        result[row, :length] = full[row, mask[row]]
    return result


def _gather_ids(full, mask, *, pad_value: int):
    lengths = mask.sum(dim=1).tolist()
    result = full.new_full((full.shape[0], max(lengths)), pad_value)
    valid = torch.zeros_like(result, dtype=torch.bool)
    for row, length in enumerate(lengths):
        result[row, :length] = full[row, mask[row]]
        valid[row, :length] = True
    return result, valid


def _replace_rows(full, mask, rows):
    result = full.clone()
    for batch_index in range(full.shape[0]):
        count = int(mask[batch_index].sum())
        result[batch_index, mask[batch_index]] = rows[batch_index, :count]
    return result


class GemmaCVRR(nn.Module):
    """One localized Gemma backbone with one shared recurrent-cell LoRA."""

    def __init__(
        self,
        model_path: str,
        *,
        ell_star: int,
        steps: int = 4,
        beta: float = 0.33,
        rank: int = 32,
        alpha: float = 12.0,
        dropout: float = 0.01,
        device: str | torch.device = "cuda:0",
        offline: bool = True,
    ):
        super().__init__()
        from transformers import AutoConfig

        self.model_path = str(model_path)
        self.device_ref = torch.device(device)
        config = AutoConfig.from_pretrained(
            model_path, local_files_only=offline
        )
        if config.model_type == "gemma3":
            from transformers import Gemma3ForConditionalGeneration as ModelClass
        elif config.model_type == "gemma4_unified":
            from transformers import (
                Gemma4UnifiedForConditionalGeneration as ModelClass,
            )
            if int(getattr(config.text_config, "num_kv_shared_layers", 0)):
                raise NotImplementedError(
                    "Gemma4 cross-layer shared KV would be an upper visual bypass"
                )
        else:
            raise ValueError(f"unsupported Gemma model_type={config.model_type!r}")
        self.base_model = ModelClass.from_pretrained(
            model_path,
            dtype=torch.bfloat16,
            device_map=str(self.device_ref),
            local_files_only=offline,
            attn_implementation="sdpa",
        )
        self.model_type = config.model_type
        self.layers = self.base_model.model.language_model.layers
        self.ell_star = int(ell_star)
        self.cell_index = self.ell_star + 1
        self.upper_start = self.cell_index + 1
        if not 0 <= self.ell_star <= len(self.layers) - 2:
            raise ValueError(
                f"ell_star={ell_star} must leave a recurrent cell and upper decoder"
            )
        if steps < 2:
            raise ValueError("CVRR training requires at least two recurrent states")
        if not 0.0 <= beta <= 1.0:
            raise ValueError("beta must lie in [0,1]")
        self.steps = int(steps)
        self.beta = float(beta)
        for parameter in self.base_model.parameters():
            parameter.requires_grad_(False)
        self.lora = inject_cell_lora(
            self.layers[self.cell_index],
            rank=rank,
            alpha=alpha,
            dropout=dropout,
        )
        # The dense checkpoint was placed before adapters were constructed;
        # newly allocated A/B tensors otherwise remain on CPU until an outer
        # training entrypoint happens to call ``model.to(device)``.
        for module in self.lora.values():
            module.to(self.device_ref)
        self.rank = int(rank)
        self.alpha = float(alpha)
        self.adapter_dropout = float(dropout)
        self.base_model.eval()

    @contextlib.contextmanager
    def adapters(self, enabled: bool):
        previous = [module.enabled for module in self.lora.values()]
        for module in self.lora.values():
            module.enabled = bool(enabled)
        try:
            yield
        finally:
            for module, value in zip(self.lora.values(), previous):
                module.enabled = value

    def train(self, mode: bool = True):
        super().train(mode)
        # Frozen dense modules stay deterministic; only LoRA dropout follows
        # training mode.
        self.base_model.eval()
        for module in self.lora.values():
            module.dropout.train(mode)
        return self

    def _modality(self, mm_inputs):
        if "token_type_ids" in mm_inputs:
            return mm_inputs["token_type_ids"]
        if "mm_token_type_ids" in mm_inputs:
            return mm_inputs["mm_token_type_ids"]
        raise ValueError("Gemma multimodal inputs have no modality IDs")

    def _pad_token_id(self) -> int:
        return int(self.base_model.config.text_config.pad_token_id)

    def _initial_multimodal(self, mm_inputs):
        calls: dict[int, LayerCall] = {}
        captured = {}
        cell = self.layers[self.cell_index]

        def stop(_module, _args, output):
            captured["hidden"] = _layer_hidden(output).detach()
            raise _StopAfterCell

        pre = cell.register_forward_pre_hook(
            _capture_call(calls, self.cell_index), with_kwargs=True
        )
        post = cell.register_forward_hook(stop)
        try:
            with torch.no_grad(), self.adapters(False):
                try:
                    self.base_model.model(
                        **mm_inputs, use_cache=False, return_dict=True
                    )
                except _StopAfterCell:
                    pass
        finally:
            pre.remove()
            post.remove()
        if "hidden" not in captured or self.cell_index not in calls:
            raise RuntimeError("failed to capture native multimodal recurrent cell")
        return captured["hidden"], calls[self.cell_index]

    def _text_context(self, question_ids, question_mask):
        calls: dict[int, LayerCall] = {}
        captured = {}
        handles = []
        for index in range(self.cell_index, len(self.layers)):
            handles.append(
                self.layers[index].register_forward_pre_hook(
                    _capture_call(calls, index), with_kwargs=True
                )
            )

        def capture_cell(_module, _args, output):
            captured["anchor"] = _layer_hidden(output).detach()

        handles.append(self.layers[self.cell_index].register_forward_hook(capture_cell))
        try:
            with torch.no_grad(), self.adapters(False):
                self.base_model.model(
                    input_ids=question_ids,
                    attention_mask=question_mask,
                    use_cache=False,
                    return_dict=True,
                )
        finally:
            for handle in handles:
                handle.remove()
        missing = [
            index
            for index in range(self.cell_index, len(self.layers))
            if index not in calls
        ]
        if missing or "anchor" not in captured:
            raise RuntimeError(f"failed to capture text context; missing={missing}")
        return captured["anchor"], calls

    @staticmethod
    def _call_layer(layer, hidden, call: LayerCall):
        return _layer_hidden(
            layer(hidden, *call.positional_tail, **call.keywords)
        )

    def _upper(self, state, text_calls):
        hidden = state
        with self.adapters(False):
            for index in range(self.upper_start, len(self.layers)):
                hidden = self._call_layer(
                    self.layers[index], hidden, text_calls[index]
                )
        hidden = self.base_model.model.language_model.norm(hidden)
        logits = self.base_model.lm_head(hidden)
        if self.model_type == "gemma4_unified":
            cap = self.base_model.config.text_config.final_logit_softcapping
            if cap is not None:
                logits = torch.tanh(logits / cap) * cap
        return logits

    def extract(self, mm_inputs: dict[str, torch.Tensor]) -> GemmaCVRRTrace:
        """Extract the frozen native first read and text-only upper context."""

        attention = mm_inputs["attention_mask"].bool()
        visual = self._modality(mm_inputs).eq(1) & attention
        text_rows = (~visual) & attention
        question_ids, question_valid = _gather_ids(
            mm_inputs["input_ids"],
            text_rows,
            pad_value=self._pad_token_id(),
        )
        question_mask = question_valid.long()
        first_full, mm_cell_call = self._initial_multimodal(mm_inputs)
        base_anchor, text_calls = self._text_context(question_ids, question_mask)
        r1 = _gather_rows(first_full, text_rows)
        if r1.shape != base_anchor.shape:
            raise RuntimeError(
                "native multimodal and image-free question states are misaligned: "
                f"R1={tuple(r1.shape)}, B={tuple(base_anchor.shape)}"
            )
        return GemmaCVRRTrace(
            scaffold=first_full.detach(),
            text_rows=text_rows,
            visual_rows=visual,
            question_valid=question_valid,
            base_anchor=base_anchor.detach(),
            r1=r1.detach(),
            mm_cell_call=mm_cell_call,
            text_calls=text_calls,
        )

    def rollout(
        self,
        trace: GemmaCVRRTrace,
        *,
        steps: int | None = None,
        initial_state: torch.Tensor | None = None,
    ) -> list[torch.Tensor]:
        """Run the shared native cell and return ``[R1, ..., R_T]``."""

        horizon = self.steps if steps is None else int(steps)
        if horizon < 1:
            raise ValueError("rollout steps must be positive")
        state = trace.r1 if initial_state is None else initial_state
        if state.shape != trace.r1.shape:
            raise ValueError(
                f"initial state shape {tuple(state.shape)} != {tuple(trace.r1.shape)}"
            )
        states = [state]
        for _ in range(1, horizon):
            recurrent_input = _replace_rows(
                trace.scaffold, trace.text_rows, state
            )
            with self.adapters(True):
                proposal_full = self._call_layer(
                    self.layers[self.cell_index],
                    recurrent_input,
                    trace.mm_cell_call,
                )
            proposal = _gather_rows(proposal_full, trace.text_rows)
            state = state + self.beta * (proposal - state)
            states.append(state)
        return states

    def decode_logits(
        self,
        trace: GemmaCVRRTrace,
        state: torch.Tensor,
    ) -> torch.Tensor:
        """Decode one question-shaped state through the strict text-only path."""

        if state.shape != trace.base_anchor.shape:
            raise ValueError(
                f"decoder state shape {tuple(state.shape)} != "
                f"text anchor {tuple(trace.base_anchor.shape)}"
            )
        # Written explicitly as B + C_T to mirror the method definition.  No
        # multimodal row or multimodal cache is passed to `_upper`.
        decoder_state = trace.base_anchor + (state - trace.base_anchor)
        return self._upper(decoder_state, trace.text_calls).float()

    def next_token_logits(
        self,
        trace: GemmaCVRRTrace,
        state: torch.Tensor,
    ) -> torch.Tensor:
        """Return the distribution after each sample's final valid prompt row."""

        logits = self.decode_logits(trace, state)
        row = trace.question_lengths.to(logits.device) - 1
        if bool((row < 0).any()):
            raise ValueError("empty question sequence")
        batch = torch.arange(logits.shape[0], device=logits.device)
        return logits[batch, row]

    def residual(self, trace: GemmaCVRRTrace, state: torch.Tensor) -> torch.Tensor:
        return state - trace.base_anchor

    def state_from_residual(
        self,
        trace: GemmaCVRRTrace,
        residual: torch.Tensor,
    ) -> torch.Tensor:
        if residual.shape != trace.base_anchor.shape:
            raise ValueError(
                f"residual shape {tuple(residual.shape)} != "
                f"text anchor {tuple(trace.base_anchor.shape)}"
            )
        return trace.base_anchor + residual

    def forward(self, mm_inputs: dict[str, torch.Tensor], mm_labels):
        question_labels, _ = _gather_ids(
            mm_labels,
            ((~self._modality(mm_inputs).eq(1)) & mm_inputs["attention_mask"].bool()),
            pad_value=-100,
        )
        trace = self.extract(mm_inputs)
        state = self.rollout(trace)[-1]
        logits = self.decode_logits(trace, state)
        shift_logits = logits[:, :-1]
        shift_labels = question_labels[:, 1:]
        token_loss = F.cross_entropy(
            shift_logits.reshape(-1, shift_logits.shape[-1]),
            shift_labels.reshape(-1),
            ignore_index=-100,
            reduction="none",
        ).reshape(shift_labels.shape)
        valid = shift_labels.ne(-100)
        counts = valid.sum(dim=1).clamp_min(1)
        per_example = (token_loss * valid).sum(dim=1) / counts
        return {
            "loss": per_example.mean(),
            "logits": logits,
            "labels": question_labels,
            "r1": trace.r1.detach(),
            "rT": state.detach(),
            "base_anchor": trace.base_anchor.detach(),
            "visual_rows": (
                self._modality(mm_inputs).eq(1)
                & mm_inputs["attention_mask"].bool()
            ).sum(dim=1).detach(),
        }

    def adapter_state_dict(self):
        return {
            name: tensor.detach().cpu()
            for name, tensor in self.state_dict().items()
            if ".lora_A" in name or ".lora_B" in name
        }

    def save_adapter(self, output_dir: str | pathlib.Path, *, step: int):
        output = pathlib.Path(output_dir)
        output.mkdir(parents=True, exist_ok=True)
        torch.save(self.adapter_state_dict(), output / "adapter_model.pt")
        metadata = {
            "format": "gemma_cvrr_lora_v1",
            "base_model": self.model_path,
            "model_type": self.model_type,
            "ell_star": self.ell_star,
            "cell_layer": self.cell_index,
            "num_workspace_steps": self.steps,
            "counterfactual_beta": self.beta,
            "adapter_rank": self.rank,
            "adapter_alpha": self.alpha,
            "adapter_dropout": self.adapter_dropout,
            "step": int(step),
            "strict_path": True,
        }
        (output / "cvrr_config.json").write_text(json.dumps(metadata, indent=2))

    def load_adapter(self, adapter_dir: str | pathlib.Path):
        """Load an adapter exactly; missing or surplus LoRA tensors are fatal."""

        adapter_dir = pathlib.Path(adapter_dir).expanduser().resolve()
        metadata = json.loads((adapter_dir / "cvrr_config.json").read_text())
        checks = {
            "model_type": self.model_type,
            "ell_star": self.ell_star,
            "cell_layer": self.cell_index,
            "num_workspace_steps": self.steps,
            "adapter_rank": self.rank,
        }
        mismatches = {
            name: (metadata.get(name), expected)
            for name, expected in checks.items()
            if metadata.get(name) != expected
        }
        if mismatches:
            raise ValueError(f"adapter metadata mismatch: {mismatches}")
        expected_keys = set(self.adapter_state_dict())
        try:
            payload = torch.load(
                adapter_dir / "adapter_model.pt",
                map_location="cpu",
                weights_only=True,
            )
        except TypeError:
            payload = torch.load(adapter_dir / "adapter_model.pt", map_location="cpu")
        actual_keys = set(payload)
        if actual_keys != expected_keys:
            raise RuntimeError(
                "adapter tensor mismatch: "
                f"missing={sorted(expected_keys - actual_keys)[:8]}, "
                f"unexpected={sorted(actual_keys - expected_keys)[:8]}"
            )
        incompatible = self.load_state_dict(payload, strict=False)
        unexpected = list(incompatible.unexpected_keys)
        missing_lora = [key for key in incompatible.missing_keys if key in expected_keys]
        if unexpected or missing_lora:
            raise RuntimeError(
                f"adapter load failed: missing={missing_lora}, unexpected={unexpected}"
            )
        return metadata


def _prompt(processor, question: str, hint: str):
    content = [
        {"type": "image"},
        {"type": "text", "text": question.strip() + str(hint)},
    ]
    return processor.apply_chat_template(
        [{"role": "user", "content": content}],
        tokenize=False,
        add_generation_prompt=True,
    )


class GemmaArrowCollator:
    def __init__(self, processor):
        self.processor = processor
        self.tokenizer = processor.tokenizer

    @staticmethod
    def _pad_1d(values, pad):
        width = max(item.shape[0] for item in values)
        result = values[0].new_full((len(values), width), pad)
        for index, item in enumerate(values):
            result[index, : item.shape[0]] = item
        return result

    def __call__(self, features):
        from PIL import Image

        examples = []
        for feature in features:
            raw = feature["image_bytes"]
            if isinstance(raw, memoryview):
                raw = raw.tobytes()
            with Image.open(io.BytesIO(raw)) as opened:
                image = opened.convert("RGB")
            prompt = _prompt(
                self.processor,
                str(feature["fixed_question"]),
                str(feature["fixed_hint"]),
            )
            prompt_item = self.processor(
                text=prompt, images=[image], return_tensors="pt"
            )
            full_item = self.processor(
                text=prompt + str(feature["fixed_answer"]).strip(),
                images=[image],
                return_tensors="pt",
            )
            prompt_ids = prompt_item["input_ids"][0]
            full_ids = full_item["input_ids"][0]
            if not torch.equal(full_ids[: prompt_ids.numel()], prompt_ids):
                raise RuntimeError("Gemma answer serialization changed the prompt prefix")
            eos = full_ids.new_tensor([self.tokenizer.eos_token_id])
            item = {name: value for name, value in full_item.items()}
            item["input_ids"] = torch.cat((full_ids, eos))
            item["attention_mask"] = torch.cat(
                (item["attention_mask"][0], torch.ones_like(eos))
            )
            modality_name = (
                "token_type_ids"
                if "token_type_ids" in item
                else "mm_token_type_ids"
            )
            item[modality_name] = torch.cat(
                (item[modality_name][0], torch.zeros_like(eos))
            )
            answer_ids = item["input_ids"][prompt_ids.numel() :]
            item["labels"] = torch.cat(
                (torch.full_like(prompt_ids, -100), answer_ids)
            )
            examples.append(item)

        sequence_names = {
            "input_ids": int(self.tokenizer.pad_token_id),
            "attention_mask": 0,
            "labels": -100,
        }
        modality_name = (
            "token_type_ids"
            if "token_type_ids" in examples[0]
            else "mm_token_type_ids"
        )
        sequence_names[modality_name] = 0
        batch = {
            name: self._pad_1d([item[name] for item in examples], pad)
            for name, pad in sequence_names.items()
        }
        for name in examples[0]:
            if name in sequence_names or name == "labels":
                continue
            values = [item[name] for item in examples]
            batch[name] = torch.cat(values, dim=0)
        labels = batch.pop("labels")
        return {"mm_inputs": batch, "mm_labels": labels}


def move_batch(batch, device):
    return {
        "mm_inputs": {
            name: value.to(device, non_blocking=True)
            for name, value in batch["mm_inputs"].items()
        },
        "mm_labels": batch["mm_labels"].to(device, non_blocking=True),
    }