hydra_gemma4.py

"""HydraGemma4: one model, many heads.

    model = HydraGemma4.from_pretrained(
        base_model="google/gemma-4-E2B-it",
        adapter_path="./colgemma4_adapter",
    )

    # Retrieval -- LoRA ON, bidirectional attention
    embeddings = model.embed(images)           # [B, N, 128]

    # Generation -- LoRA OFF, causal attention
    text = model.generate(image, "Describe this document.")

One ColGemma4 + LoRA + lm_head. Attention mode and LoRA toggled per call.

Gemma 4 E2B layer structure:
    35 layers total: 28 sliding_attention + 7 full_attention
    full_attention at indices: [4, 9, 14, 19, 24, 29, 34]
    Only full_attention layers are patched for bidirectional mode.
    Sliding layers stay causal always.

Key Gemma 4 differences from Qwen3.5:
    - Layer forward signature includes `per_layer_input` (Per-Layer Embeddings)
    - Layer type via config.text_config.layer_types list (not layer.layer_type attribute)
    - Vision tower at self.vision_tower (not self.visual)
    - final_logit_softcapping for logit capping
    - image_position_ids (not image_grid_thw)
"""

import re
from typing import List, Union

import torch
import torch.nn as nn
from PIL import Image
from peft import PeftModel
from transformers import AutoProcessor, Gemma4ForConditionalGeneration
from transformers.models.gemma4.modeling_gemma4 import Gemma4Model as _Gemma4Base

from colgemma4 import ColGemma4, ColGemma4Processor


class HydraGemma4:
    """One ColGemma4 model with LoRA + lm_head. Two faces via toggling.

    Retrieval face:  LoRA ON,  bidirectional attention -> custom_text_proj -> 128-dim
    Generation face: LoRA OFF, causal attention -> lm_head -> autoregressive text
    """

    def __init__(self, model, lm_head, emb_processor, gen_processor, attn_fns, config):
        self.model = model
        self.lm_head = lm_head
        self.emb_processor = emb_processor
        self.gen_processor = gen_processor
        self._attn_fns = attn_fns  # {layer_idx: (causal_fn, bidir_fn)}
        self._config = config
        self._base = self._unwrap(model)

    @staticmethod
    def _unwrap(model):
        m = model
        if hasattr(m, "module"):
            m = m.module
        if hasattr(m, "base_model"):
            m = m.base_model
        if hasattr(m, "model"):
            m = m.model
        return m

    def _set_bidirectional(self):
        for idx, (_, bidir_fn) in self._attn_fns.items():
            self._base.language_model.layers[idx].forward = bidir_fn

    def _set_causal(self):
        for idx, (causal_fn, _) in self._attn_fns.items():
            self._base.language_model.layers[idx].forward = causal_fn

    @classmethod
    def from_pretrained(
        cls,
        base_model: str = "google/gemma-4-E2B-it",
        adapter_path: str = "/tmp/colgemma4_adapter",
        torch_dtype=torch.bfloat16,
        max_visual_tokens: int = 560,
        device: str = "cuda",
    ):
        """Load single model from base + adapter + lm_head."""
        from pathlib import Path

        print(f"Loading ColGemma4 from {base_model}...", flush=True)
        config = Gemma4ForConditionalGeneration.config_class.from_pretrained(
            base_model, trust_remote_code=True
        )
        config.text_config.use_cache = False

        model = ColGemma4.from_pretrained(
            base_model,
            config=config,
            torch_dtype=torch_dtype,
            attn_implementation="sdpa",
            ignore_mismatched_sizes=True,
        )

        # Store causal/bidirectional forwards for full-attention layers
        layer_types = config.text_config.layer_types
        attn_fns = {}
        for idx, layer in enumerate(model.language_model.layers):
            if layer_types[idx] == "full_attention":
                causal_fn = layer.forward

                def make_bidir(orig):
                    def bidir(
                        hidden_states,
                        per_layer_input=None,
                        position_embeddings=None,
                        attention_mask=None,
                        **kw,
                    ):
                        if (
                            attention_mask is not None
                            and attention_mask.ndim == 4
                            and attention_mask.dtype.is_floating_point
                        ):
                            min_dtype = torch.finfo(attention_mask.dtype).min
                            diag = torch.diagonal(attention_mask, dim1=-2, dim2=-1)
                            is_valid = diag > (min_dtype / 2)
                            bidir_mask = torch.where(
                                is_valid.unsqueeze(-1) & is_valid.unsqueeze(-2),
                                attention_mask.new_zeros(1),
                                attention_mask.new_full((1,), min_dtype),
                            )
                            attention_mask = bidir_mask
                        return orig(
                            hidden_states,
                            per_layer_input=per_layer_input,
                            position_embeddings=position_embeddings,
                            attention_mask=attention_mask,
                            **kw,
                        )

                    return bidir

                bidir_fn = make_bidir(causal_fn)
                attn_fns[idx] = (causal_fn, bidir_fn)
                layer.forward = bidir_fn

        print(f"Patched {len(attn_fns)} full-attention layers", flush=True)

        print(f"Loading LoRA adapter from {adapter_path}...", flush=True)
        model = PeftModel.from_pretrained(model, adapter_path)
        model = model.to(device).eval()

        # Load lm_head
        print("Loading lm_head...", flush=True)
        lm_head_path = Path(adapter_path) / "lm_head.pt"
        if lm_head_path.exists():
            base_cfg = Gemma4ForConditionalGeneration.config_class.from_pretrained(base_model)
            lm_head = nn.Linear(
                base_cfg.text_config.hidden_size,
                base_cfg.text_config.vocab_size,
                bias=False,
            )
            lm_head.load_state_dict(torch.load(lm_head_path, map_location="cpu"))
        else:
            base = Gemma4ForConditionalGeneration.from_pretrained(
                base_model, torch_dtype=torch_dtype
            )
            lm_head = base.lm_head
            del base
            torch.cuda.empty_cache()

        lm_head = lm_head.to(device).to(torch_dtype)

        emb_processor = ColGemma4Processor.from_pretrained(
            base_model, max_num_visual_tokens=max_visual_tokens
        )
        gen_processor = AutoProcessor.from_pretrained(base_model, trust_remote_code=True)

        params = sum(p.numel() for p in model.parameters()) / 1e6
        lm_params = sum(p.numel() for p in lm_head.parameters()) / 1e6
        print(
            f"Ready. Single model: {params:.0f}M + lm_head: {lm_params:.0f}M",
            flush=True,
        )

        return cls(model, lm_head, emb_processor, gen_processor, attn_fns, config)

    def embed(self, images: Union[Image.Image, List[Image.Image]]) -> torch.Tensor:
        """Embed images for retrieval. LoRA ON, bidirectional attention."""
        if isinstance(images, Image.Image):
            images = [images]
        self.model.enable_adapter_layers()
        self._set_bidirectional()
        inputs = self.emb_processor.process_images(images)
        device = next(self.model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
            return self.model(**inputs)

    def embed_queries(self, queries: Union[str, List[str]]) -> torch.Tensor:
        """Embed queries for retrieval. LoRA ON, bidirectional attention."""
        if isinstance(queries, str):
            queries = [queries]
        self.model.enable_adapter_layers()
        self._set_bidirectional()
        inputs = self.emb_processor.process_queries(queries)
        device = next(self.model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
            return self.model(**inputs)

    @torch.no_grad()
    def generate(
        self,
        image: Image.Image,
        prompt: str,
        max_new_tokens: int = 4096,
        system_prompt: str = None,
        temperature: float = 1.0,
        top_p: float = 0.95,
        top_k: int = 64,
    ) -> str:
        """Generate text given image + prompt. LoRA OFF, causal attention."""
        self.model.disable_adapter_layers()
        self._set_causal()
        try:
            return self._generate_with_kv_cache(
                image, prompt, max_new_tokens, system_prompt,
                temperature=temperature, top_p=top_p, top_k=top_k,
            )
        except Exception:
            return self._generate_no_cache(
                image, prompt, max_new_tokens, system_prompt,
                temperature=temperature, top_p=top_p, top_k=top_k,
            )
        finally:
            self.model.enable_adapter_layers()
            self._set_bidirectional()

    def _prepare_inputs(self, image, prompt, system_prompt=None):
        msgs = []
        if system_prompt:
            msgs.append({"role": "system", "content": system_prompt})
        msgs.append({
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt},
            ],
        })
        txt = self.gen_processor.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True,
        )
        inp = self.gen_processor(
            text=[txt], images=[image], return_tensors="pt", padding=True
        )
        device = next(self.model.parameters()).device
        return {k: v.to(device) for k, v in inp.items()}

    def _decode(self, generated):
        if not generated:
            return ""
        gen_ids = torch.cat(generated, dim=1)
        return self.gen_processor.tokenizer.decode(gen_ids[0], skip_special_tokens=True).strip()

    def _is_eos(self, token_id):
        eos = self.gen_processor.tokenizer.eos_token_id
        if isinstance(eos, list):
            return token_id in eos
        return token_id == eos

    @staticmethod
    def _sample_token(logits, temperature, top_p, top_k):
        if temperature <= 0:
            return logits.argmax(dim=-1)
        logits_1d = logits[0, 0] / temperature
        if top_k > 0:
            topk_vals, _ = torch.topk(logits_1d, min(top_k, logits_1d.size(-1)))
            logits_1d[logits_1d < topk_vals[-1]] = float("-inf")
        probs = torch.softmax(logits_1d, dim=-1)
        if 0 < top_p < 1.0:
            sorted_probs, sorted_idx = torch.sort(probs, descending=True)
            cumsum = torch.cumsum(sorted_probs, dim=-1)
            mask = cumsum - sorted_probs > top_p
            sorted_probs[mask] = 0
            probs = torch.zeros_like(probs).scatter_(0, sorted_idx, sorted_probs)
        probs = probs / probs.sum()
        token_id = torch.multinomial(probs, num_samples=1)
        return token_id.reshape(1, 1)

    def _generate_with_kv_cache(
        self, image, prompt, max_new_tokens, system_prompt=None,
        temperature=1.0, top_p=0.95, top_k=64,
    ):
        """KV-cache generation. Pixel values only on first step."""
        inp = self._prepare_inputs(image, prompt, system_prompt)
        input_ids = inp["input_ids"]
        attn_mask = inp["attention_mask"]
        mm_ids = inp.get("mm_token_type_ids")
        past_key_values = None
        generated = []

        final_softcap = self._config.text_config.final_logit_softcapping

        for step in range(max_new_tokens):
            if step == 0:
                kw = {
                    "input_ids": input_ids,
                    "attention_mask": attn_mask,
                    "pixel_values": inp.get("pixel_values"),
                    "image_position_ids": inp.get("image_position_ids"),
                    "use_cache": True,
                    "output_hidden_states": True,
                    "return_dict": True,
                }
                if mm_ids is not None:
                    kw["mm_token_type_ids"] = mm_ids
            else:
                kw = {
                    "input_ids": next_token,
                    "attention_mask": attn_mask,
                    "past_key_values": past_key_values,
                    "use_cache": True,
                    "output_hidden_states": True,
                    "return_dict": True,
                }

            outputs = _Gemma4Base.forward(self._base, **kw)
            past_key_values = outputs.past_key_values
            logits = self.lm_head(outputs.last_hidden_state[:, -1:, :])

            if final_softcap is not None:
                logits = logits / final_softcap
                logits = torch.tanh(logits)
                logits = logits * final_softcap

            next_token = self._sample_token(logits, temperature, top_p, top_k)
            generated.append(next_token)

            if self._is_eos(next_token.item()):
                break

            attn_mask = torch.cat([attn_mask, torch.ones_like(next_token)], dim=1)

        return self._decode(generated)

    def _generate_no_cache(
        self, image, prompt, max_new_tokens, system_prompt=None,
        temperature=1.0, top_p=0.95, top_k=64,
    ):
        """No-cache fallback. Passes all inputs every step."""
        inp = self._prepare_inputs(image, prompt, system_prompt)
        input_ids = inp["input_ids"]
        attn_mask = inp["attention_mask"]
        mm_ids = inp.get("mm_token_type_ids")
        generated = []

        final_softcap = self._config.text_config.final_logit_softcapping

        for step in range(max_new_tokens):
            kw = {
                "input_ids": input_ids,
                "attention_mask": attn_mask,
                "pixel_values": inp.get("pixel_values"),
                "image_position_ids": inp.get("image_position_ids"),
            }
            if mm_ids is not None:
                kw["mm_token_type_ids"] = mm_ids

            hidden = _Gemma4Base.forward(
                self._base, **kw,
                use_cache=False, output_hidden_states=True, return_dict=True,
            ).last_hidden_state

            logits = self.lm_head(hidden[:, -1:, :])
            if final_softcap is not None:
                logits = logits / final_softcap
                logits = torch.tanh(logits)
                logits = logits * final_softcap

            next_token = self._sample_token(logits, temperature, top_p, top_k)
            generated.append(next_token)

            if self._is_eos(next_token.item()):
                break

            input_ids = torch.cat([input_ids, next_token], dim=1)
            attn_mask = torch.cat([attn_mask, torch.ones_like(next_token)], dim=1)
            if mm_ids is not None:
                mm_ids = torch.cat([mm_ids, torch.zeros_like(next_token)], dim=1)

        return self._decode(generated)