File size: 3,522 Bytes
62d02e9 a22274b 62d02e9 a22274b 62d02e9 a22274b 62d02e9 a22274b 62d02e9 a22274b 62d02e9 a22274b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Serve unlabelled Kev 0.5B typed decisions from the Core ML artifact."""
from __future__ import annotations
import argparse
import json
import shutil
import tempfile
from pathlib import Path
import coremltools as ct
import numpy as np
from kev.api import output_tokens, to_answers
from transformers import AutoTokenizer
from preprocessing import Shape, prepare_runtime_inputs
COMPUTE_UNITS = {
"all": ct.ComputeUnit.ALL,
"cpu": ct.ComputeUnit.CPU_ONLY,
"cpu-gpu": ct.ComputeUnit.CPU_AND_GPU,
"cpu-ne": ct.ComputeUnit.CPU_AND_NE,
}
PACKAGE = "kev_0_5b_fp16_L128_options32.mlpackage"
def stage_symlinked_package(package: Path) -> tuple[Path, tempfile.TemporaryDirectory | None]:
"""Give Core ML real package files when a Hub snapshot contains blob symlinks."""
if not package.is_symlink() and not any(path.is_symlink() for path in package.rglob("*")):
return package, None
temporary = tempfile.TemporaryDirectory(prefix="kev-coreml-")
staged = Path(temporary.name) / package.name
shutil.copytree(package, staged, symlinks=False)
return staged, temporary
class KevCoreML:
"""Reusable serving session; initialization and request latency stay separate."""
def __init__(self, model_dir: str | Path, *, units: str = "all", package: str | Path | None = None):
root = Path(model_dir)
self.tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer", local_files_only=True)
self.shape = Shape()
model_path = Path(package) if package is not None else root / PACKAGE
model_path, self._staged_package = stage_symlinked_package(model_path)
self.model = ct.models.MLModel(str(model_path), compute_units=COMPUTE_UNITS[units])
def predict(self, request: dict) -> dict:
arrays, encoded, metadata, parsed = prepare_runtime_inputs(self.tokenizer, request, self.shape)
result = np.asarray(self.model.predict(arrays)["probabilities"], dtype=np.float64)
count = len(metadata[0]["keys"])
if result.shape != (1, self.shape.max_options) or not np.isfinite(result).all():
raise ValueError("Invalid Core ML probabilities")
selected = result[0, :count].tolist()
keys = metadata[0]["keys"]
answers = to_answers([selected], metadata)
return {
"model": parsed.model,
"answers": answers,
"usage": {
"input_tokens": len(encoded["ids"]),
"output_tokens": output_tokens(self.tokenizer, answers),
},
"option_keys": keys,
"probabilities": selected,
"chosen": keys[int(np.argmax(selected))],
}
def predict(model_dir: str | Path, request: dict, *, units: str = "all", package: str | Path | None = None) -> dict:
"""One-shot compatibility wrapper; reuse `KevCoreML` for repeated calls."""
return KevCoreML(model_dir, units=units, package=package).predict(request)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model-dir", required=True, type=Path)
parser.add_argument("--request-json", type=Path, required=True)
parser.add_argument("--units", choices=COMPUTE_UNITS, default="all")
parser.add_argument("--package", type=Path)
args = parser.parse_args()
request = json.loads(args.request_json.read_text())
print(json.dumps(predict(args.model_dir, request, units=args.units, package=args.package), indent=2))
if __name__ == "__main__":
main()
|