Spaces:
Running
Running
| """MLX backend for Apple silicon: the same readout from an mlx-lm model on the Mac. One forward pass per question; with | |
| `prefix_cache=True` the shared 'State:' prefix of a request is prefilled once and copied for the other questions of that step. | |
| MLX is not thread-safe: all model calls go through one lock in a worker thread, so the async server stays responsive. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import copy | |
| import threading | |
| from typing import Any | |
| class MlxBackend: | |
| def __init__(self, model_dir: str, *, prefix_cache: bool = True, min_prefix_tokens: int = 64, letter_prefix: str = "auto", assistant_prefix: str = ""): | |
| import mlx.core as mx | |
| from mlx_lm import load | |
| from transformers import AutoTokenizer | |
| self.mx = mx | |
| self.model, _ = load(model_dir) | |
| self.tok = AutoTokenizer.from_pretrained(model_dir) | |
| self.prefix_cache = prefix_cache | |
| self.min_prefix_tokens = min_prefix_tokens | |
| self._lock = threading.Lock() | |
| self._prefix: tuple[list[int], Any] | None = None | |
| from .letters import header_counter, letter_ids, render | |
| self._ids, self.letter_prefix = letter_ids(self.tok, letter_prefix) | |
| self.assistant_prefix = assistant_prefix | |
| self._render = lambda content: render(self.tok, content, assistant_prefix) | |
| self._header = header_counter(self.tok, assistant_prefix) | |
| self.vision = False # mlx-lm text models only | |
| async def start(self) -> dict: | |
| return { | |
| "served_model": "mlx", | |
| "exact_readout": True, | |
| "chat_template": bool(getattr(self.tok, "chat_template", None)), | |
| "letter_prefix": self.letter_prefix, | |
| "vision": False, | |
| "prefix_cache": self.prefix_cache, | |
| } | |
| def letter_ids(self) -> dict[str, int]: | |
| return self._ids | |
| def header_tokens(self, text: str) -> int: | |
| return self._header(text) | |
| async def healthy(self) -> bool: | |
| return True | |
| async def aclose(self): | |
| pass | |
| def _ids_for(self, content) -> list[int]: | |
| if not isinstance(content, str): | |
| raise TypeError("the MLX build is text-only: screenshot states are not supported") | |
| return list(self.tok.encode(self._render(content), add_special_tokens=False)) | |
| def _feed(self, ids: list[int], cache, step: int = 2048): | |
| mx = self.mx | |
| while len(ids) > step: | |
| self.model(mx.array(ids[:step])[None], cache=cache) | |
| mx.eval([c.state for c in cache]) | |
| ids = ids[step:] | |
| return self.model(mx.array(ids)[None], cache=cache)[0, -1] | |
| def _last_logits(self, ids: list[int], content: str): | |
| mx = self.mx | |
| cut = content.rfind("\n\nQuestion: ") if self.prefix_cache else -1 | |
| if cut < 0: | |
| return self.model(mx.array(ids)[None])[0, -1] | |
| from mlx_lm.models.cache import make_prompt_cache | |
| shared = self._prefix[0] if self._prefix else None | |
| if not (shared and ids[: len(shared)] == shared and len(ids) > len(shared)): | |
| head = self._ids_for(content[:cut]) | |
| n = next((i for i, (x, y) in enumerate(zip(ids, head, strict=False)) if x != y), min(len(ids), len(head))) | |
| n = min(n, len(ids) - 1) | |
| if n < self.min_prefix_tokens: | |
| return self.model(mx.array(ids)[None])[0, -1] | |
| cache = make_prompt_cache(self.model) | |
| self._feed(ids[:n], cache) | |
| mx.eval([c.state for c in cache]) | |
| self._prefix = (ids[:n], cache) | |
| shared = ids[:n] | |
| return self._feed(ids[len(shared) :], copy.deepcopy(self._prefix[1])) | |
| def _logprobs_sync(self, content: str, token_ids: list[int]) -> tuple[list[float], int]: | |
| mx = self.mx | |
| ids = self._ids_for(content) | |
| with self._lock: | |
| logits = self._last_logits(ids, content).astype(mx.float32) | |
| lp = logits - mx.logsumexp(logits, axis=-1) | |
| vals = [float(lp[i]) for i in token_ids] | |
| mx.eval(logits) | |
| return vals, len(ids) | |
| async def logprobs(self, content: str | list[dict[str, Any]], token_ids: list[int]) -> tuple[list[float], int]: | |
| return await asyncio.to_thread(self._logprobs_sync, content, token_ids) | |
| async def prewarm(self, text: str) -> int: | |
| def run(): | |
| ids = self._ids_for(text) | |
| with self._lock: | |
| if self.prefix_cache: | |
| from mlx_lm.models.cache import make_prompt_cache | |
| cache = make_prompt_cache(self.model) | |
| self._feed(ids, cache) | |
| self.mx.eval([c.state for c in cache]) | |
| self._prefix = (ids, cache) | |
| return len(ids) | |
| return await asyncio.to_thread(run) | |