File size: 5,320 Bytes
307b208 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | 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))
|