| """Use NanoJev's pinned native renderer for decision requests.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import numpy as np | |
| from assets import upstream_module | |
| def prepare_request(root: Path, tokenizer, request: dict, length: int, max_candidates: int): | |
| native = upstream_module(root, "predict_toy_decisions") | |
| examples = native.prepare_examples(request, tokenizer, length) | |
| if len(examples) != 1: | |
| raise ValueError("This Core ML bucket accepts one question per call") | |
| example = examples[0] | |
| paths = example["leaf_tokens"] | |
| if len(paths) > max_candidates: | |
| raise ValueError(f"{len(paths)} candidates exceed bucket capacity {max_candidates}") | |
| pad = tokenizer.pad_token_id | |
| ids = np.full((max_candidates, length), pad, dtype=np.int32) | |
| attention = np.zeros((max_candidates, length), dtype=np.int32) | |
| eos_map = np.zeros((max_candidates, 1, length), dtype=np.float32) | |
| candidate_mask = np.zeros((1, max_candidates), dtype=np.float32) | |
| for index, tokens in enumerate(paths): | |
| ids[index, : len(tokens)] = tokens | |
| attention[index, : len(tokens)] = 1 | |
| eos_map[index, 0, len(tokens) - 1] = 1 | |
| candidate_mask[0, index] = 1 | |
| return {"input_ids": ids, "attention_mask": attention, "eos_map": eos_map}, candidate_mask, example | |