File size: 6,724 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 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """EmbeddingGemma 2 adapter. Neural computation runs exclusively in Core AI."""
import json
from pathlib import Path
import numpy as np
from coreai.runtime import AIModel, NDArray, SpecializationOptions
from transformers import AutoProcessor
class EmbeddingGemma2CoreAI:
@classmethod
async def load(cls, bundle):
self = cls()
self.root = Path(bundle)
self.manifest = json.loads((self.root / "manifest.json").read_text())
self.processor = AutoProcessor.from_pretrained(self.root)
self.prompts = json.loads(
(self.root / "config_sentence_transformers.json").read_text()
)["prompts"]
self.models = {}
self.functions = {}
for name, asset in [
("text", "embedding"),
("vision", "vision"),
("audio", "audio"),
]:
self.models[name] = await AIModel.load(
self.root / f"{asset}.aimodel", SpecializationOptions.default()
)
self.functions[name] = self.models[name].load_function("main")
return self
async def call(self, name, inputs):
result = await self.functions[name](
{k: NDArray(np.ascontiguousarray(v)) for k, v in inputs.items()}
)
return [result[k].numpy() for k in self.manifest[f"{name}_outputs"]]
async def encode(
self,
text=None,
*,
images=None,
videos=None,
audio=None,
sampling_rate=16000,
prompt_name=None,
truncate_dim=768,
):
if not 1 <= truncate_dim <= 768:
raise ValueError("truncate_dim must be between 1 and 768")
if prompt_name:
if images is not None or videos is not None or audio is not None:
raise ValueError("Task prompts apply to text-only inputs")
text = self.prompts[prompt_name] + (text or "")
kwargs = {"return_tensors": "np"}
if text is not None:
kwargs["text"] = text
if images is not None:
kwargs["images"] = images
if videos is not None:
if not isinstance(videos, np.ndarray) or videos.ndim != 4:
raise ValueError(
"Video must be an already sampled NumPy [frames,height,width,channels] array"
)
kwargs["videos"] = videos
kwargs["videos_kwargs"] = {"do_sample_frames": False}
if audio is not None:
samples = np.asarray(audio)
if sampling_rate != 16000:
raise ValueError("Audio must be sampled at 16000 Hz")
if samples.ndim != 1 or len(samples) > 179200:
raise ValueError("One mono clip of at most 11.2 seconds is supported")
kwargs.update(audio=samples, sampling_rate=sampling_rate)
encoded = self.processor(**kwargs)
return await self.encode_processed(encoded, truncate_dim=truncate_dim)
async def encode_processed(self, e, truncate_dim=768):
ids = np.asarray(e["input_ids"], dtype=np.int32)
if ids.shape[0] != 1:
raise ValueError(
"Batch size must be one; call encode separately per example"
)
n = ids.shape[1]
cap = self.manifest["sequence_length"]
if n > cap:
raise ValueError(
f"Input needs {n} tokens, capacity {cap}; input was not truncated"
)
media = np.zeros((1, cap, 512), np.float32)
media_mask = np.zeros((1, cap, 1), bool)
p = self.processor
for pixels_key, pos_key, token in [
("pixel_values", "image_position_ids", p.image_token_id),
("pixel_values_videos", "video_position_ids", p.video_token_id),
]:
if pixels_key not in e:
continue
features = []
for pixels, positions in zip(e[pixels_key], e[pos_key]):
count = len(pixels)
capacity = self.manifest["vision_shape"][1]
if count > capacity:
raise ValueError("Vision patch capacity exceeded")
pix = np.pad(
np.asarray(pixels, dtype=np.float32),
((0, capacity - count), (0, 0)),
)
pos = np.pad(
np.asarray(positions, dtype=np.int32),
((0, capacity - count), (0, 0)),
constant_values=-1,
)
values, valid = await self.call(
"vision",
{"pixel_values": pix[None], "pixel_position_ids": pos[None]},
)
features.append(values[valid.astype(bool)])
features = np.concatenate(features, axis=0)
slots = np.flatnonzero(ids[0] == token)
if len(slots) != len(features):
raise ValueError("Visual placeholder count differs from encoder output")
media[0, slots] = features
media_mask[0, slots] = True
if "input_features" in e:
features = []
for mel, mask in zip(e["input_features"], e["input_features_mask"]):
count = len(mel)
capacity = self.manifest["audio_shape"][1]
if count > capacity:
raise ValueError("Audio feature capacity exceeded")
mel = np.pad(
np.asarray(mel, dtype=np.float32), ((0, capacity - count), (0, 0))
)
mask = np.pad(np.asarray(mask, dtype=bool), (0, capacity - count))
values, valid = await self.call(
"audio",
{"input_features": mel[None], "input_features_mask": mask[None]},
)
features.append(values[valid.astype(bool)])
features = np.concatenate(features, axis=0)
slots = np.flatnonzero(ids[0] == p.audio_token_id)
if len(slots) != len(features):
raise ValueError("Audio placeholder count differs from encoder output")
media[0, slots] = features
media_mask[0, slots] = True
padded = np.pad(ids, ((0, 0), (0, cap - n)))
mask = np.pad(
np.asarray(e["attention_mask"], dtype=np.int32), ((0, 0), (0, cap - n))
)
value = (
await self.call(
"text",
{
"input_ids": padded,
"attention_mask": mask,
"media_embeddings": media,
"media_mask": media_mask,
},
)
)[0][0, :truncate_dim]
return value / np.linalg.norm(value)
|