| import json |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from coreai_torch import TorchConverter, get_decomp_table |
| from transformers import AutoModel, AutoProcessor |
|
|
|
|
| def unfold(x, dimension, size, step): |
| dimension %= x.ndim |
| return torch.stack( |
| [ |
| x.narrow(dimension, i, size) |
| for i in range(0, x.shape[dimension] - size + 1, step) |
| ], |
| dim=dimension, |
| ).movedim(dimension + 1, -1) |
|
|
|
|
| def save(model, sample, path): |
| with torch.no_grad(): |
| ep = torch.export.export(model, sample, strict=False) |
| table = get_decomp_table() |
| table[torch.ops.aten.unfold.default] = unfold |
| ep = ep.run_decompositions(table) |
| print("Converting", path, flush=True) |
| TorchConverter().add_exported_program(ep).to_coreai().save_asset(path) |
| return list(ep.graph_signature.user_outputs) |
|
|
|
|
| class Text(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.m = m.language_model |
|
|
| def forward(self, input_ids, attention_mask, media_embeddings, media_mask): |
| emb = self.m.embed_tokens(input_ids) |
| emb = torch.where(media_mask, media_embeddings.to(emb.dtype), emb) |
| n = input_ids.shape[1] |
| pos = torch.arange(n, device=input_ids.device).unsqueeze(0) |
| mask = attention_mask[:, None, None, :].bool().expand(1, 1, n, n) |
| sliding = ( |
| pos[:, :, None] - pos[:, None, :] |
| ).abs() <= self.m.config.sliding_window |
| hidden = self.m( |
| inputs_embeds=emb, |
| position_ids=pos, |
| attention_mask={ |
| "full_attention": mask, |
| "sliding_attention": mask & sliding[:, None], |
| }, |
| ).last_hidden_state |
| weights = attention_mask[:, :, None].float() |
| pooled = (hidden.float() * weights).sum(1) / weights.sum(1).clamp(min=1) |
| return F.normalize(pooled, dim=-1) |
|
|
|
|
| class Vision(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.v = m.vision_tower |
| self.project = m.embed_vision |
|
|
| def forward(self, pixel_values, pixel_position_ids): |
| v = self.v |
| pad = (pixel_position_ids == -1).all(-1) |
| emb = v.patch_embedder(pixel_values, pixel_position_ids, pad) |
| hidden = v.encoder( |
| inputs_embeds=emb, |
| attention_mask=~pad, |
| pixel_position_ids=pixel_position_ids, |
| ).last_hidden_state |
| hidden, mask = v.pooler( |
| hidden, pixel_position_ids, pad, pixel_values.shape[1] // 9 |
| ) |
| return self.project(hidden.to(emb.dtype)), mask |
|
|
|
|
| class Audio(torch.nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.a = m.audio_tower |
| self.project = m.embed_audio |
|
|
| def forward(self, input_features, input_features_mask): |
| out = self.a(input_features, attention_mask=input_features_mask) |
| return self.project(out.last_hidden_state), out.attention_mask |
|
|
|
|
| if __name__ == "__main__": |
| torch.set_num_threads(4) |
| root = Path("bundles/embeddinggemma-2-CoreAI") |
| root.mkdir(parents=True, exist_ok=True) |
| m = AutoModel.from_pretrained( |
| "upstream/embeddinggemma-2", |
| dtype=torch.float32, |
| attn_implementation={ |
| "language_model": "sdpa", |
| "vision_tower": "sdpa", |
| "audio_tower": "eager", |
| }, |
| ).eval() |
| print("Loaded", flush=True) |
| manifest = { |
| "source": "google/embeddinggemma-2", |
| "revision": "914f7f89142e33e77833254d9c9b90c3cef7303b", |
| "precision": "float32", |
| "sequence_length": 1024, |
| "transformers_revision": "536ecc007387a50e77603bb5d92100e9b07514cc", |
| } |
| sample = ( |
| torch.ones(1, 1024, dtype=torch.int32), |
| torch.ones(1, 1024, dtype=torch.int32), |
| torch.zeros(1, 1024, 512), |
| torch.zeros(1, 1024, 1, dtype=torch.bool), |
| ) |
| manifest["text_outputs"] = save(Text(m).eval(), sample, root / "embedding.aimodel") |
| (root / "manifest.json").write_text(json.dumps(manifest, indent=2)) |
| p = AutoProcessor.from_pretrained("upstream/embeddinggemma-2") |
| from PIL import Image |
|
|
| e = p(images=Image.new("RGB", (224, 224)), return_tensors="pt") |
| print("Image shapes", [(k, tuple(v.shape)) for k, v in e.items()], flush=True) |
| sample = (e["pixel_values"].float(), e["image_position_ids"].int()) |
| manifest["vision_outputs"] = save(Vision(m).eval(), sample, root / "vision.aimodel") |
| manifest["vision_shape"] = list(sample[0].shape) |
| (root / "manifest.json").write_text(json.dumps(manifest, indent=2)) |
| import numpy as np |
|
|
| e = p( |
| audio=np.sin(np.arange(16000) * 0.1).astype(np.float32), |
| sampling_rate=16000, |
| return_tensors="pt", |
| ) |
| print("Audio shapes", [(k, tuple(v.shape)) for k, v in e.items()], flush=True) |
| sample = ( |
| F.pad( |
| e["input_features"].float(), (0, 0, 0, 1120 - e["input_features"].shape[1]) |
| ), |
| F.pad( |
| e["input_features_mask"].bool(), |
| (0, 1120 - e["input_features_mask"].shape[1]), |
| ), |
| ) |
| manifest["audio_outputs"] = save(Audio(m).eval(), sample, root / "audio.aimodel") |
| manifest["audio_shape"] = list(sample[0].shape) |
| (root / "manifest.json").write_text(json.dumps(manifest, indent=2)) |
|
|
