| """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) |
|
|
