# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""The EmbeddingGemma 2 text encoder with multimodal (soft-token) inputs: `eg2_text.EmbeddingGemma2Text` plus a
mix of token embeddings and externally computed soft tokens at the start of the graph.

    input_ids      [1, S] int32   token ids; placeholder slots (<|image|> = 258880) are looked up like any id
    attention_mask [1, S] fp16    1 = token, 0 = right padding
    mm_embeds      [1, S, 512] fp16  soft tokens placed at their slots (zeros elsewhere); NOT scaled by sqrt(512)
    mm_mask        [1, S] fp16    1 at a soft-token slot, 0 elsewhere
    -> x = embed(input_ids) * (1 - mm_mask) + mm_embeds * mm_mask, then exactly the text graph (per-layer
       embeddings, 24 layers, final norm, projection, masked mean pooling over ALL tokens, L2 normalization)
    -> embedding [1, 768] fp16

With mm_mask all zero this is the text encoder (x * 1 + 0 is exact in fp16). eg2_text.py is unchanged (it is the
file that produced the text-only bundle); this module only re-states its forward with the mix.
"""
from __future__ import annotations

import torch

import eg2_text


class EmbeddingGemma2TextMM(eg2_text.EmbeddingGemma2Text):
    def forward(self, input_ids, attention_mask, mm_embeds, mm_mask):  # noqa: D102
        S = input_ids.shape[1]
        keep = (self.one - mm_mask)[:, :, None]
        x = self.embed_tokens(input_ids) * keep + mm_embeds * mm_mask[:, :, None]  # [1, S, 512]
        ple = self.per_layer_projection_norm(
            self.per_layer_model_projection(x).view(1, S, self.cfg.n_layers, self.cfg.ple_dim))
        full, sliding = self._biases(attention_mask, S)
        cs = {}
        for lt in ("sliding_attention", "full_attention"):
            cs[lt] = (getattr(self, f"cos_{lt}")[:S].to(x.dtype)[None, None], getattr(self, f"sin_{lt}")[:S].to(x.dtype)[None, None])
        for i, layer in enumerate(self.layers):
            lt = self.cfg.layer_types[i]
            x = layer(x, ple[:, :, i, :], cs[lt][0], cs[lt][1], sliding if lt == "sliding_attention" else full)
        x = self.embedding_projection(self.norm(x))
        if eg2_text.PRECISION_MODE == "ane":
            w = attention_mask * eg2_text._recip(attention_mask.sum(-1, keepdim=True))
            pooled = torch.matmul(w[:, None, :], x)[:, 0, :]
            u = pooled * eg2_text._recip(torch.maximum(pooled.abs().amax(-1, keepdim=True), self.floor16))
            return u * torch.rsqrt((u * u).sum(-1, keepdim=True))
        xf = x.float()
        wf = attention_mask.float()
        pooled = (xf * wf[:, :, None]).sum(1) / wf.sum(-1, keepdim=True)
        return (pooled / pooled.norm(dim=-1, keepdim=True)).to(x.dtype)


def load_model(model_dir: str, dtype: torch.dtype, max_seq: int = 2048) -> EmbeddingGemma2TextMM:
    """Same weights and folding as eg2_text.load_model, mm forward."""
    cfg = eg2_text.Cfg(model_dir)
    m = EmbeddingGemma2TextMM(cfg, max_seq)
    base = eg2_text.load_model(model_dir, dtype, max_seq)
    m.load_state_dict(base.state_dict(), strict=True)
    return m.to(dtype).eval()