| from __future__ import annotations |
|
|
| import time |
|
|
| from .base import BaseProvider, ProviderResponse |
|
|
|
|
| class MLXProvider(BaseProvider): |
| name = "mlx" |
|
|
| def __init__(self, model_name: str | None, adapter_path: str | None = None, max_tokens: int = 400): |
| if not model_name: |
| raise ValueError("MLX provider requires --model.") |
| from mlx_lm import generate, load |
| from mlx_lm.sample_utils import make_sampler |
|
|
| self._generate = generate |
| self._sampler = make_sampler(temp=0.0) |
| self._model, self._tokenizer = load(model_name, adapter_path=adapter_path) |
| self.model_name = model_name |
| self.adapter_path = adapter_path |
| self.max_tokens = max_tokens |
|
|
| def generate(self, card: dict, messages: list[dict[str, str]], prompt_text: str) -> ProviderResponse: |
| prompt = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| start = time.perf_counter() |
| raw_text = self._generate( |
| self._model, |
| self._tokenizer, |
| prompt=prompt, |
| max_tokens=self.max_tokens, |
| verbose=False, |
| sampler=self._sampler, |
| ) |
| latency_ms = (time.perf_counter() - start) * 1000 |
| return ProviderResponse( |
| raw_text=raw_text.strip(), |
| latency_ms=latency_ms, |
| metadata={"model": self.model_name, "adapter": self.adapter_path}, |
| ) |
|
|
