"""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) # Qwen3.5 mRoPE: text increments all axes, media contributes its grid coordinates. 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}