| |
| """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): |
| S = input_ids.shape[1] |
| keep = (self.one - mm_mask)[:, :, None] |
| x = self.embed_tokens(input_ids) * keep + mm_embeds * mm_mask[:, :, None] |
| 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() |
|
|
