| """Serve the pinned System One contract from the published Core ML packages.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import os |
| from pathlib import Path |
|
|
| import torch |
| from packing import MarkerCollator |
| from transformers import AutoTokenizer |
| from typed_coreml import KINDS, marker_tensor, padded_batch |
| from upstream_system_one import _answer, system_one_records |
|
|
|
|
| class CoreMLSystemOne: |
| def __init__(self, repo_root: Path, embedding_w8_kinds: tuple[str, ...] = ()): |
| self.root = repo_root.resolve() |
| if not set(embedding_w8_kinds).issubset(KINDS): |
| raise ValueError(f"unknown compressed decision path: {embedding_w8_kinds}") |
| self.embedding_w8_kinds = frozenset(embedding_w8_kinds) |
| lock = json.loads((self.root / "conversion" / "assets.lock.json").read_text()) |
| self.model_name = { |
| "decision-kai": "Decision-1.0-Kai", |
| "decision-lex": "Decision-1.0-Lex", |
| }[lock["engine"]] |
| tokenizer = AutoTokenizer.from_pretrained( |
| self.root / "tokenizer", local_files_only=True, trust_remote_code=False |
| ) |
| |
| self.collator = MarkerCollator(tokenizer, max_length=1024, state_truncation="error") |
| self.models = {} |
| self.model_inputs = {} |
|
|
| def _model(self, kind: str): |
| import coremltools as ct |
|
|
| if kind not in self.models: |
| |
| self.models.clear() |
| self.model_inputs.clear() |
| gc.collect() |
| suffix = "-embedding-w8" if kind in self.embedding_w8_kinds else "" |
| path = self.root / "coreml" / (kind + suffix + ".mlpackage") |
| if not path.is_dir(): |
| raise FileNotFoundError(f"selected Core ML package is missing: {path}") |
| self.models[kind] = ct.models.MLModel( |
| os.path.relpath(path, Path.cwd()), compute_units=ct.ComputeUnit.CPU_AND_NE |
| ) |
| return self.models[kind] |
|
|
| def predict_row(self, record: dict) -> dict: |
| kind = record["question"]["type"].lower() |
| metadata = json.loads((self.root / "reports" / (kind + ".json")).read_text()) |
| shape = metadata["shape"] |
| if shape["batch"] != 1: |
| raise ValueError("This package expects one question per model call") |
| batch, encoded = padded_batch( |
| self.collator, [record], shape["tokens"], shape["candidates"] |
| ) |
| model = self._model(kind) |
| inputs = { |
| "input_ids": batch["input_ids"].to(torch.int32).numpy(), |
| "attention_mask": batch["attention_mask"].to(torch.int32).numpy(), |
| } |
| if kind not in self.model_inputs: |
| self.model_inputs[kind] = {item.name for item in model.get_spec().description.input} |
| input_names = self.model_inputs[kind] |
| if "marker_map" in input_names: |
| inputs["marker_map"] = marker_tensor(batch, marker_map=True).numpy() |
| elif "marker_positions" in input_names: |
| inputs["marker_positions"] = marker_tensor(batch, marker_map=False).numpy() |
| else: |
| raise ValueError("Core ML package lacks a supported marker input") |
| result = model.predict(inputs) |
| count = len(encoded[0]["positions"]) |
| logits = torch.from_numpy(result["logits"][0, :count]).float() |
| probabilities = logits.softmax(-1) |
| row = encoded[0] |
| output = { |
| key: row[key] |
| for key in ( |
| "id", "question_id", "candidate_ids", "state_tokens_original", |
| "state_tokens_kept", "input_tokens" |
| ) |
| } |
| output.update( |
| type=kind.capitalize(), |
| logits=logits.tolist(), |
| probabilities=probabilities.tolist(), |
| confidence=float(probabilities.max()), |
| ) |
| if kind == "noul": |
| output["probability"] = float(probabilities[1]) |
| else: |
| output["choice_id"] = row["candidate_ids"][int(probabilities.argmax())] |
| if kind == "score": |
| output["score"] = float( |
| (probabilities * torch.arange(count, dtype=probabilities.dtype)).sum() |
| ) |
| output["expected_value"] = float( |
| (probabilities * batch["values"][0, :count]).sum() |
| ) |
| return output |
|
|
| def evaluate(self, request: dict) -> dict: |
| if request.get("model") != self.model_name: |
| raise ValueError("Request model does not match the published checkpoint") |
| rows = system_one_records(request) |
| answers = {} |
| input_tokens = 0 |
| for row in rows: |
| prediction = self.predict_row(row) |
| answers[row["question"]["id"]] = _answer(row, prediction) |
| input_tokens += prediction["input_tokens"] |
| return { |
| "model": self.model_name, |
| "answers": answers, |
| "usage": {"input_tokens": input_tokens, "output_tokens": 0}, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--repo-root", required=True, type=Path) |
| parser.add_argument("--request-json", required=True, type=Path) |
| parser.add_argument("--embedding-w8-kinds", nargs="+", choices=KINDS, default=(), |
| help="use optional embedding-W8 packages for the listed typed paths") |
| args = parser.parse_args() |
| runtime = CoreMLSystemOne(args.repo_root, tuple(args.embedding_w8_kinds)) |
| result = runtime.evaluate(json.loads(args.request_json.read_text())) |
| print(json.dumps(result, indent=2, ensure_ascii=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
