| """Decision 3.0 host adapter. All neural inference runs through Core AI.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| from coreai.runtime import AIModel, NDArray, SpecializationOptions |
| from PIL import Image |
| from transformers import AutoProcessor, AutoTokenizer |
|
|
| from d3_format import SYSTEM_PROMPT, user_prompt |
| from d3_format import render as render_d3 |
| from d3_runtime import load_image, load_video, normalize_question, product_answer |
|
|
|
|
| def pack(encoded, manifest, features=None): |
| ids = np.asarray(encoded["input_ids"], dtype=np.int32) |
| n = ids.shape[1] |
| length = manifest["sequence_length"] |
| hidden = manifest["hidden_size"] |
| if n > length: |
| raise ValueError( |
| f"Prompt needs {n} tokens; asset capacity is {length}; input was not truncated" |
| ) |
| pad = length - n |
| ids = np.pad(ids, ((0, 0), (pad, 0)), constant_values=0) |
| mask = np.pad(np.ones((1, n), dtype=np.int32), ((0, 0), (pad, 0))) |
| pos = np.zeros((4, 1, length), dtype=np.int32) |
| pos[0, 0, pad:] = np.arange(n) |
| media_mask = np.zeros((1, length, 1), dtype=np.bool_) |
| media = np.zeros((1, length, hidden), dtype=np.float32) |
| |
| types = np.asarray( |
| encoded.get("mm_token_type_ids", np.zeros((1, n))), dtype=np.int32 |
| )[0] |
| cursor = 0 |
| offset = 0 |
| feature_index = 0 |
| grid_side = manifest["image_side"] // 16 // 2 |
| while cursor < n: |
| if types[cursor] == 0: |
| end = cursor + 1 |
| while end < n and types[end] == 0: |
| end += 1 |
| pos[1:, 0, pad + cursor : pad + end] = np.arange( |
| offset, offset + end - cursor |
| ) |
| offset += end - cursor |
| cursor = end |
| else: |
| count = grid_side**2 |
| end = cursor + count |
| if end > n or not np.all(types[cursor:end] == types[cursor]): |
| raise ValueError("Unexpected visual grid") |
| pos[1, 0, pad + cursor : pad + end] = offset |
| pos[2, 0, pad + cursor : pad + end] = offset + np.repeat( |
| np.arange(grid_side), grid_side |
| ) |
| pos[3, 0, pad + cursor : pad + end] = offset + np.tile( |
| np.arange(grid_side), grid_side |
| ) |
| media_mask[0, pad + cursor : pad + end] = True |
| media[0, pad + cursor : pad + end] = features[feature_index] |
| feature_index += 1 |
| offset += grid_side |
| cursor = end |
| return dict( |
| input_ids=ids, |
| attention_mask=mask, |
| position_ids=pos, |
| visual_embeddings=media, |
| visual_mask=media_mask, |
| ) |
|
|
|
|
| class D3CoreAI: |
| @classmethod |
| async def load(cls, bundle, cpu_only=False): |
| self = cls() |
| self.root = Path(bundle) |
| self.manifest = json.loads((self.root / "manifest.json").read_text()) |
| self.config = json.loads((self.root / "decision_config.json").read_text()) |
| self.tokenizer = AutoTokenizer.from_pretrained(self.root) |
| self.processor = AutoProcessor.from_pretrained(self.root) |
| opts = ( |
| SpecializationOptions.cpu_only() |
| if cpu_only |
| else SpecializationOptions.default() |
| ) |
| self.model = await AIModel.load(self.root / "decision.aimodel", opts) |
| self.fn = self.model.load_function("main") |
| self.vision = await AIModel.load(self.root / "vision.aimodel", opts) |
| self.vfn = self.vision.load_function("main") |
| return self |
|
|
| async def logits(self, state, question, images=None, videos=None): |
| normalized = normalize_question(question) |
| rendered, keys = normalized.rendered, normalized.keys |
| features = [] |
| if images or videos: |
| side = self.manifest["image_side"] |
| resized = [ |
| load_image(x).resize((side, side), Image.Resampling.BICUBIC) |
| for x in (images or []) |
| ] |
| decoded = [load_video(x) for x in (videos or [])] |
| frames = [ |
| np.stack( |
| [ |
| np.asarray( |
| Image.fromarray(frame).resize( |
| (side, side), Image.Resampling.BICUBIC |
| ) |
| ) |
| for frame in v.frames |
| ] |
| ) |
| for v in decoded |
| ] |
| messages = [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| { |
| "role": "user", |
| "content": [{"type": "image"} for _ in resized] |
| + [{"type": "video"} for _ in decoded] |
| + [ |
| { |
| "type": "text", |
| "text": user_prompt(state, rendered, self.config["codes"]), |
| } |
| ], |
| }, |
| ] |
| text = self.processor.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| args = { |
| "text": [text], |
| "return_tensors": "np", |
| "images_kwargs": {"do_resize": False}, |
| "videos_kwargs": { |
| "do_resize": False, |
| "do_sample_frames": False, |
| "return_metadata": True, |
| }, |
| } |
| if resized: |
| args["images"] = resized |
| if decoded: |
| args.update( |
| videos=frames, video_metadata=[v.metadata() for v in decoded] |
| ) |
| encoded = self.processor(**args) |
| if len(encoded["input_ids"][0]) > self.manifest["sequence_length"]: |
| raise ValueError( |
| "Media prompt exceeds asset capacity; input was not truncated" |
| ) |
| per_image = (side // 16) ** 2 |
| for name in ["pixel_values", "pixel_values_videos"]: |
| if name not in encoded: |
| continue |
| patches = encoded[name] |
| for start in range(0, len(patches), per_image): |
| out = await self.vfn( |
| { |
| "pixel_values": NDArray( |
| patches[start : start + per_image].astype(np.float32) |
| ) |
| } |
| ) |
| features.append(out[self.manifest["vision_output"]].numpy()) |
| else: |
| text = render_d3(self.tokenizer, state, rendered, self.config["codes"]) |
| encoded = self.tokenizer( |
| [text], add_special_tokens=False, return_tensors="np" |
| ) |
| arrays = pack(encoded, self.manifest, features) |
| out = await self.fn({k: NDArray(v) for k, v in arrays.items()}) |
| return out[self.manifest["output"]].numpy()[0, : len(keys)], keys, encoded |
|
|
| async def system_one(self, *, state, questions, images=None, videos=None): |
| answers = {} |
| for key, q in questions.items(): |
| logits, keys, _ = await self.logits(state, q, images, videos) |
| z = logits / self.config.get("temperature", 1.0) |
| probs = np.exp(z - z.max()) |
| probs /= probs.sum() |
| if not np.isfinite(probs).all(): |
| raise FloatingPointError("Core AI returned nonfinite probabilities") |
| answers[key] = product_answer(normalize_question(q), probs.tolist()) |
| return {"model": self.manifest["source"], "answers": answers} |
|
|
