llm-coreml 0.1.0__tar.gz

llm_coreml-0.1.0/PKG-INFO

@@ -0,0 +1,190 @@
+Metadata-Version: 2.3
+Name: llm-coreml
+Version: 0.1.0
+Summary: A plugin for https://github.com/simonw/llm for running local CoreML .mlpackage model files.
+Author: Anentropic
+Author-email: Anentropic <ego@anentropic.com>
+Requires-Dist: llm
+Requires-Dist: coremltools>=8.0
+Requires-Dist: transformers
+Requires-Dist: numpy
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# llm-coreml
+
+A plugin for <https://llm.datasette.io/> cli tool that runs CoreML `.mlpackage` LLM models locally on macOS.
+
+Point it at a model (and its corresponding HuggingFace tokenizer), then prompt it like any other `llm` model.
+
+## Requirements
+
+- macOS (CoreML is Apple-only)
+- Python 3.11+
+
+## Installation
+
+```bash
+llm install llm-coreml
+```
+
+Or for development:
+
+```bash
+git clone https://github.com/anentropic/llm-coreml.git
+cd llm-coreml
+llm install -e .
+```
+
+## Quick start
+
+Register a model with a name and a path to the `.mlpackage`.
+
+The `--tokenizer` argument is the HuggingFace model name to load the tokenizer from. This should match the HF model your `.mlpackage` was derived from:
+
+```bash
+llm coreml add my-llama /path/to/llama.mlpackage \
+    --tokenizer meta-llama/Llama-3.2-1B-Instruct
+```
+
+Prompt it:
+
+```bash
+llm -m coreml/my-llama "Explain quantum computing in one sentence"
+```
+
+Check it shows up in `llm models`:
+
+```bash
+llm models | grep coreml
+```
+
+## Usage
+
+### Prompting
+
+```bash
+# Basic prompt
+llm -m coreml/my-llama "What is Rust?"
+
+# With a system prompt
+llm -m coreml/my-llama "Hello" -s "You are a pirate"
+
+# Continue a conversation
+llm -m coreml/my-llama "What is Rust?"
+llm -c "Compare it to Go"
+```
+
+### Model options
+
+Pass options with `-o`:
+
+```bash
+llm -m coreml/my-llama "Write a haiku" \
+    -o temperature 0.7 \
+    -o top_p 0.9 \
+    -o max_tokens 50
+```
+
+### Python API
+
+```python
+import llm
+
+model = llm.get_model("coreml/my-llama")
+response = model.prompt("What is the capital of France?")
+print(response.text())
+```
+
+## CLI reference
+
+### `llm coreml add`
+
+```
+llm coreml add <name> <path> --tokenizer <hf_id> [--compute-units <units>]
+```
+
+Register a CoreML model.
+
+| Argument | Description |
+|---|---|
+| `name` | Model name, used as `coreml/<name>` |
+| `path` | Path to the `.mlpackage` directory (resolved to absolute) |
+| `--tokenizer` | HuggingFace tokenizer model ID (required) |
+| `--compute-units` | Compute units: `all`, `cpu_only`, `cpu_and_gpu`, `cpu_and_ne` (default: `all`) |
+
+### `llm coreml list`
+
+```
+llm coreml list
+```
+
+Lists registered models with their paths, tokenizer IDs, and compute units.
+
+### `llm coreml remove`
+
+```
+llm coreml remove <name>
+```
+
+Removes a registered model. Exits with code 1 if the model doesn't exist.
+
+## Model options reference
+
+| Option | Type | Default | Description |
+|---|---|---|---|
+| `max_tokens` | int | 200 | Maximum tokens to generate |
+| `temperature` | float | 0.0 | Sampling temperature. 0 = greedy (deterministic) |
+| `top_p` | float | 1.0 | Top-p nucleus sampling threshold |
+
+## How it works
+
+### Format auto-detection
+
+The plugin reads the CoreML model spec at load time and checks the input names:
+
+- `inputIds` (camelCase) = Apple format, uses float16 causal masks
+- `input_ids` (snake_case) = HuggingFace format, uses int32 attention masks
+
+No config file needed.
+
+### Stateful KV-cache
+
+If the model spec declares `stateDescriptions`, the plugin uses stateful inference with KV-cache. Otherwise it falls back to stateless inference, which reprocesses the full sequence each step (slower, but works with older models).
+
+### Tokenization
+
+The plugin uses `transformers.AutoTokenizer` with `apply_chat_template()` to handle chat formatting. The tokenizer is downloaded and cached the first time you use a model.
+
+## Getting CoreML models
+
+You can get `.mlpackage` LLM models by:
+
+- Converting HuggingFace models with [coremltools](https://apple.github.io/coremltools/docs-guides/)
+- Using Apple's [ml-explore](https://github.com/ml-explore) tools
+- Downloading pre-converted models from HuggingFace (search for "coreml" tagged models)
+
+## Development
+
+```bash
+uv sync --dev
+```
+
+### Quality gates
+
+```bash
+uv run basedpyright      # Type checking (strict)
+uv run ruff check        # Linting
+uv run ruff format       # Formatting
+uv run pytest            # Tests
+```
+
+Or all at once:
+
+```bash
+prek run --all-files
+```
+
+## License
+
+MIT

llm_coreml-0.1.0/README.md

@@ -0,0 +1,177 @@
+# llm-coreml
+
+A plugin for <https://llm.datasette.io/> cli tool that runs CoreML `.mlpackage` LLM models locally on macOS.
+
+Point it at a model (and its corresponding HuggingFace tokenizer), then prompt it like any other `llm` model.
+
+## Requirements
+
+- macOS (CoreML is Apple-only)
+- Python 3.11+
+
+## Installation
+
+```bash
+llm install llm-coreml
+```
+
+Or for development:
+
+```bash
+git clone https://github.com/anentropic/llm-coreml.git
+cd llm-coreml
+llm install -e .
+```
+
+## Quick start
+
+Register a model with a name and a path to the `.mlpackage`.
+
+The `--tokenizer` argument is the HuggingFace model name to load the tokenizer from. This should match the HF model your `.mlpackage` was derived from:
+
+```bash
+llm coreml add my-llama /path/to/llama.mlpackage \
+    --tokenizer meta-llama/Llama-3.2-1B-Instruct
+```
+
+Prompt it:
+
+```bash
+llm -m coreml/my-llama "Explain quantum computing in one sentence"
+```
+
+Check it shows up in `llm models`:
+
+```bash
+llm models | grep coreml
+```
+
+## Usage
+
+### Prompting
+
+```bash
+# Basic prompt
+llm -m coreml/my-llama "What is Rust?"
+
+# With a system prompt
+llm -m coreml/my-llama "Hello" -s "You are a pirate"
+
+# Continue a conversation
+llm -m coreml/my-llama "What is Rust?"
+llm -c "Compare it to Go"
+```
+
+### Model options
+
+Pass options with `-o`:
+
+```bash
+llm -m coreml/my-llama "Write a haiku" \
+    -o temperature 0.7 \
+    -o top_p 0.9 \
+    -o max_tokens 50
+```
+
+### Python API
+
+```python
+import llm
+
+model = llm.get_model("coreml/my-llama")
+response = model.prompt("What is the capital of France?")
+print(response.text())
+```
+
+## CLI reference
+
+### `llm coreml add`
+
+```
+llm coreml add <name> <path> --tokenizer <hf_id> [--compute-units <units>]
+```
+
+Register a CoreML model.
+
+| Argument | Description |
+|---|---|
+| `name` | Model name, used as `coreml/<name>` |
+| `path` | Path to the `.mlpackage` directory (resolved to absolute) |
+| `--tokenizer` | HuggingFace tokenizer model ID (required) |
+| `--compute-units` | Compute units: `all`, `cpu_only`, `cpu_and_gpu`, `cpu_and_ne` (default: `all`) |
+
+### `llm coreml list`
+
+```
+llm coreml list
+```
+
+Lists registered models with their paths, tokenizer IDs, and compute units.
+
+### `llm coreml remove`
+
+```
+llm coreml remove <name>
+```
+
+Removes a registered model. Exits with code 1 if the model doesn't exist.
+
+## Model options reference
+
+| Option | Type | Default | Description |
+|---|---|---|---|
+| `max_tokens` | int | 200 | Maximum tokens to generate |
+| `temperature` | float | 0.0 | Sampling temperature. 0 = greedy (deterministic) |
+| `top_p` | float | 1.0 | Top-p nucleus sampling threshold |
+
+## How it works
+
+### Format auto-detection
+
+The plugin reads the CoreML model spec at load time and checks the input names:
+
+- `inputIds` (camelCase) = Apple format, uses float16 causal masks
+- `input_ids` (snake_case) = HuggingFace format, uses int32 attention masks
+
+No config file needed.
+
+### Stateful KV-cache
+
+If the model spec declares `stateDescriptions`, the plugin uses stateful inference with KV-cache. Otherwise it falls back to stateless inference, which reprocesses the full sequence each step (slower, but works with older models).
+
+### Tokenization
+
+The plugin uses `transformers.AutoTokenizer` with `apply_chat_template()` to handle chat formatting. The tokenizer is downloaded and cached the first time you use a model.
+
+## Getting CoreML models
+
+You can get `.mlpackage` LLM models by:
+
+- Converting HuggingFace models with [coremltools](https://apple.github.io/coremltools/docs-guides/)
+- Using Apple's [ml-explore](https://github.com/ml-explore) tools
+- Downloading pre-converted models from HuggingFace (search for "coreml" tagged models)
+
+## Development
+
+```bash
+uv sync --dev
+```
+
+### Quality gates
+
+```bash
+uv run basedpyright      # Type checking (strict)
+uv run ruff check        # Linting
+uv run ruff format       # Formatting
+uv run pytest            # Tests
+```
+
+Or all at once:
+
+```bash
+prek run --all-files
+```
+
+## License
+
+MIT

llm_coreml-0.1.0/pyproject.toml

@@ -0,0 +1,60 @@
+[project]
+name = "llm-coreml"
+version = "0.1.0"
+description = "A plugin for https://github.com/simonw/llm for running local CoreML .mlpackage model files."
+readme = "README.md"
+authors = [
+    { name = "Anentropic", email = "ego@anentropic.com" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "llm",
+    "coremltools>=8.0",
+    "transformers",
+    "numpy",
+]
+
+[project.entry-points.llm]
+llm_coreml = "llm_coreml"
+
+[build-system]
+requires = ["uv_build>=0.9.18,<1.0.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "basedpyright>=1.38.0",
+    "ipython>=9.10.0",
+    "pdbpp>=0.12.0.post1",
+    "pytest>=8.0.0",
+    "pytest-cov>=6.0.0",
+    "ruff>=0.15.1",
+]
+[tool.basedpyright]
+pythonVersion = "3.14"
+typeCheckingMode = "strict"
+include = ["src", "tests"]
+reportPrivateUsage = false
+reportMissingTypeStubs = false
+
+[tool.ruff]
+target-version = "py311"
+line-length = 100
+src = ["src"]
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "UP", "B", "SIM", "TCH", "D"]
+ignore = [
+    "D1",    # don't require docstrings everywhere
+    "D202",  # allow blank line after docstring (common in tests)
+    "D203",  # conflicts with D211 (no blank line before class docstring)
+    "D212",  # conflicts with D213 (multiline summary on second line)
+    "D401",  # imperative mood is awkward for dunder methods
+    "D413",  # blank line after last section not required
+]
+
+[tool.pytest.ini_options]
+markers = [
+    "integration: requires a real CoreML model (run 'just convert-test-model' first)",
+]
+addopts = "-v"

llm_coreml-0.1.0/src/llm_coreml/__init__.py

@@ -0,0 +1,83 @@
+"""A plugin for llm that runs local CoreML .mlpackage model files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import click
+import llm
+
+from llm_coreml.registry import add_model, list_models, remove_model
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+@llm.hookimpl
+def register_models(register: Callable[..., object]) -> None:
+    """Register all CoreML models from the local registry."""
+    from llm_coreml.model import CoreMLModel
+
+    for name, config in list_models().items():
+        register(
+            CoreMLModel(
+                model_id=f"coreml/{name}",
+                model_path=config["path"],
+                tokenizer_id=config["tokenizer"],
+                compute_units=config.get("compute_units", "all"),
+            ),
+        )
+
+
+@llm.hookimpl
+def register_commands(cli: click.Group) -> None:
+    """Register the `llm coreml` command group."""
+
+    @cli.group(name="coreml")
+    def coreml_group() -> None:
+        """Manage locally-registered CoreML models."""
+
+    @coreml_group.command(name="add")
+    @click.argument("name")
+    @click.argument("path", type=click.Path(exists=True))
+    @click.option("--tokenizer", required=True, help="HuggingFace tokenizer model ID")
+    @click.option(
+        "--compute-units",
+        type=click.Choice(["all", "cpu_only", "cpu_and_gpu", "cpu_and_ne"]),
+        default="all",
+        help="CoreML compute units to use (default: all)",
+    )
+    def add_cmd(name: str, path: str, tokenizer: str, compute_units: str) -> None:  # pyright: ignore[reportUnusedFunction]
+        """
+        Register a CoreML model.
+
+        NAME is the model name (used as coreml/NAME).
+        PATH is the path to the .mlpackage directory.
+        """
+        resolved = str(Path(path).resolve())
+        add_model(name, resolved, tokenizer, compute_units)
+        click.echo(f"Added model coreml/{name}")
+
+    @coreml_group.command(name="list")
+    def list_cmd() -> None:  # pyright: ignore[reportUnusedFunction]
+        """List registered CoreML models."""
+        models = list_models()
+        if not models:
+            click.echo("No CoreML models registered.")
+            return
+        for name, config in models.items():
+            click.echo(f"coreml/{name}")
+            click.echo(f"  Path:           {config['path']}")
+            click.echo(f"  Tokenizer:      {config['tokenizer']}")
+            click.echo(f"  Compute units:  {config.get('compute_units', 'all')}")
+
+    @coreml_group.command(name="remove")
+    @click.argument("name")
+    def remove_cmd(name: str) -> None:  # pyright: ignore[reportUnusedFunction]
+        """Remove a registered CoreML model."""
+        if remove_model(name):
+            click.echo(f"Removed model coreml/{name}")
+        else:
+            click.echo(f"Model '{name}' not found.", err=True)
+            raise SystemExit(1)

llm_coreml-0.1.0/src/llm_coreml/inference.py

@@ -0,0 +1,251 @@
+"""CoreML inference engine with format auto-detection and autoregressive generation."""
+
+from __future__ import annotations
+
+import enum
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+class ModelFormat(enum.Enum):
+    """Detected input format of a CoreML model."""
+
+    APPLE = "apple"
+    HUGGINGFACE = "huggingface"
+
+
+def detect_format(spec: Any) -> ModelFormat:
+    """Detect whether a model uses Apple (camelCase) or HuggingFace (snake_case) inputs."""
+    input_names: set[str] = {inp.name for inp in spec.description.input}
+    if "inputIds" in input_names:
+        return ModelFormat.APPLE
+    if "input_ids" in input_names:
+        return ModelFormat.HUGGINGFACE
+    msg = f"Cannot detect model format. Input names: {input_names}"
+    raise ValueError(msg)
+
+
+def is_stateful(spec: Any) -> bool:
+    """Check if the model supports stateful KV-cache inference."""
+    state_descs = getattr(spec.description, "stateDescriptions", None)
+    return state_descs is not None and len(state_descs) > 0
+
+
+class CoreMLInferenceEngine:
+    """Loads a CoreML model and runs autoregressive text generation."""
+
+    COMPUTE_UNIT_MAP: dict[str, str] = {
+        "all": "ALL",
+        "cpu_only": "CPU_ONLY",
+        "cpu_and_gpu": "CPU_AND_GPU",
+        "cpu_and_ne": "CPU_AND_NE",
+    }
+
+    def __init__(self, model_path: str, compute_units: str = "all") -> None:
+        self.model_path = model_path
+        self.compute_units = compute_units
+        self._model: Any = None
+        self._format: ModelFormat | None = None
+        self._stateful: bool | None = None
+
+    def _load(self) -> Any:
+        if self._model is not None:
+            return self._model
+
+        import coremltools as ct  # pyright: ignore[reportMissingTypeStubs]
+
+        cu_name = self.COMPUTE_UNIT_MAP.get(self.compute_units, "ALL")
+        cu = getattr(ct.ComputeUnit, cu_name)
+        self._model = ct.models.MLModel(
+            self.model_path,
+            compute_units=cu,
+        )
+        spec = self._model.get_spec()
+        self._format = detect_format(spec)
+        self._stateful = is_stateful(spec)
+        return self._model
+
+    @property
+    def format(self) -> ModelFormat:
+        """Model input format (loads model if needed)."""
+        self._load()
+        assert self._format is not None
+        return self._format
+
+    @property
+    def stateful(self) -> bool:
+        """Whether the model supports stateful KV-cache."""
+        self._load()
+        assert self._stateful is not None
+        return self._stateful
+
+    def generate(
+        self,
+        input_ids: list[int],
+        *,
+        max_tokens: int = 200,
+        temperature: float = 0.0,
+        top_p: float = 1.0,
+        eos_token_id: int | None = None,
+    ) -> Iterator[int]:
+        """Generate tokens autoregressively, yielding one token ID at a time."""
+        model = self._load()
+
+        if self._stateful:
+            yield from self._generate_stateful(
+                model,
+                input_ids,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                eos_token_id=eos_token_id,
+            )
+        else:
+            yield from self._generate_stateless(
+                model,
+                input_ids,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                eos_token_id=eos_token_id,
+            )
+
+    def _generate_stateful(
+        self,
+        model: Any,
+        input_ids: list[int],
+        *,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        eos_token_id: int | None,
+    ) -> Iterator[int]:
+        state = model.make_state()
+
+        # Prefill: process all prompt tokens
+        output: dict[str, np.ndarray] = {}
+        for i, token_id in enumerate(input_ids):
+            feed = self._make_input(
+                [token_id],
+                seq_len=len(input_ids),
+                position=i,
+            )
+            output = model.predict(feed, state=state)
+
+        # Decode
+        next_token = self._extract_next_token(output, temperature, top_p)
+        for _ in range(max_tokens):
+            if eos_token_id is not None and next_token == eos_token_id:
+                break
+            yield next_token
+
+            feed = self._make_input(
+                [next_token],
+                seq_len=len(input_ids) + 1,
+                position=len(input_ids),
+            )
+            input_ids.append(next_token)
+            output = model.predict(feed, state=state)
+            next_token = self._extract_next_token(output, temperature, top_p)
+
+    def _generate_stateless(
+        self,
+        model: Any,
+        input_ids: list[int],
+        *,
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        eos_token_id: int | None,
+    ) -> Iterator[int]:
+        seq = list(input_ids)
+        for _ in range(max_tokens):
+            feed = self._make_input(seq, seq_len=len(seq), position=0)
+            output: dict[str, np.ndarray] = model.predict(feed)
+            next_token = self._extract_next_token(output, temperature, top_p)
+            if eos_token_id is not None and next_token == eos_token_id:
+                break
+            yield next_token
+            seq.append(next_token)
+
+    def _make_input(
+        self,
+        token_ids: list[int],
+        *,
+        seq_len: int,
+        position: int,
+    ) -> dict[str, np.ndarray]:
+        assert self._format is not None
+        ids = np.array([token_ids], dtype=np.int32)
+
+        if self._format == ModelFormat.APPLE:
+            return {
+                "inputIds": ids,
+                "causalMask": _build_causal_mask(len(token_ids), seq_len),
+            }
+        return {
+            "input_ids": ids,
+            "attention_mask": _build_attention_mask(seq_len),
+        }
+
+    def _extract_next_token(
+        self,
+        output: dict[str, np.ndarray],
+        temperature: float,
+        top_p: float,
+    ) -> int:
+        logits_key = "logits" if "logits" in output else next(iter(output))
+        logits = np.array(output[logits_key], dtype=np.float32)
+        # Take last token's logits
+        if logits.ndim == 3:
+            logits = logits[0, -1, :]
+        elif logits.ndim == 2:
+            logits = logits[-1, :]
+        return sample_token(logits, temperature=temperature, top_p=top_p)
+
+
+def _build_causal_mask(query_len: int, kv_len: int) -> np.ndarray:
+    """Build a float16 causal mask for Apple-format models."""
+    mask = np.full((1, 1, query_len, kv_len), -np.inf, dtype=np.float16)
+    for i in range(query_len):
+        pos = kv_len - query_len + i
+        mask[0, 0, i, : pos + 1] = 0.0
+    return mask
+
+
+def _build_attention_mask(seq_len: int) -> np.ndarray:
+    """Build an int32 attention mask for HuggingFace-format models."""
+    return np.ones((1, seq_len), dtype=np.int32)
+
+
+def sample_token(
+    logits: np.ndarray,
+    *,
+    temperature: float = 0.0,
+    top_p: float = 1.0,
+) -> int:
+    """Sample a token from logits using temperature and top-p nucleus sampling."""
+    if temperature <= 0:
+        return int(np.argmax(logits))
+
+    logits = logits / temperature
+    # Numerical stability
+    logits = logits - np.max(logits)
+    probs = np.exp(logits) / np.sum(np.exp(logits))
+
+    if top_p < 1.0:
+        sorted_indices = np.argsort(-probs)
+        sorted_probs = probs[sorted_indices]
+        cumulative = np.cumsum(sorted_probs)
+        cutoff = np.searchsorted(cumulative, top_p) + 1
+        # Zero out tokens beyond cutoff
+        mask = np.zeros_like(probs)
+        mask[sorted_indices[:cutoff]] = 1.0
+        probs = probs * mask
+        probs = probs / np.sum(probs)
+
+    return int(np.random.choice(len(probs), p=probs))

llm_coreml-0.1.0/src/llm_coreml/model.py

@@ -0,0 +1,117 @@
+"""CoreML model implementation for llm."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import llm
+from pydantic import Field
+
+from llm_coreml.inference import CoreMLInferenceEngine
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+class CoreMLOptions(llm.Options):
+    """Options for CoreML model inference."""
+
+    max_tokens: int | None = Field(
+        description="Maximum number of tokens to generate",
+        default=200,
+    )
+    temperature: float | None = Field(
+        description="Sampling temperature (0 for greedy)",
+        default=0.0,
+    )
+    top_p: float | None = Field(
+        description="Top-p nucleus sampling threshold",
+        default=1.0,
+    )
+
+
+class CoreMLModel(llm.Model):
+    """An llm Model backed by a local CoreML .mlpackage."""
+
+    can_stream = True
+    model_id: str
+    model_path: str
+    tokenizer_id: str
+    Options = CoreMLOptions  # type: ignore[assignment]
+
+    def __init__(
+        self,
+        model_id: str,
+        model_path: str,
+        tokenizer_id: str,
+        compute_units: str = "all",
+    ) -> None:
+        self.model_id = model_id
+        self.model_path = model_path
+        self.tokenizer_id = tokenizer_id
+        self.compute_units = compute_units
+        self._engine: CoreMLInferenceEngine | None = None
+        self._tokenizer: Any = None
+
+    def _get_engine(self) -> CoreMLInferenceEngine:
+        if self._engine is None:
+            self._engine = CoreMLInferenceEngine(self.model_path, self.compute_units)
+        return self._engine
+
+    def _get_tokenizer(self) -> Any:
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer  # pyright: ignore[reportMissingTypeStubs]
+
+            self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_id)  # pyright: ignore[reportUnknownMemberType]
+        return self._tokenizer  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
+
+    def execute(
+        self,
+        prompt: llm.Prompt,
+        stream: bool,
+        response: llm.Response,
+        conversation: llm.Conversation | None,
+    ) -> Iterator[str]:
+        tokenizer = self._get_tokenizer()
+        engine = self._get_engine()
+
+        messages = _build_messages(prompt, conversation)
+        input_ids: list[int] = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+        )
+
+        opts: CoreMLOptions = prompt.options  # type: ignore[assignment]
+        token_count = 0
+        for token_id in engine.generate(
+            input_ids,
+            max_tokens=opts.max_tokens or 200,
+            temperature=opts.temperature or 0.0,
+            top_p=opts.top_p or 1.0,
+            eos_token_id=tokenizer.eos_token_id,
+        ):
+            token_count += 1
+            yield tokenizer.decode([token_id])  # type: ignore[no-any-return]
+
+        response.set_usage(input=len(input_ids), output=token_count)  # pyright: ignore[reportUnknownMemberType]
+
+
+def _build_messages(
+    prompt: llm.Prompt,
+    conversation: llm.Conversation | None,
+) -> list[dict[str, str]]:
+    """Reconstruct a chat message list from the conversation history."""
+    messages: list[dict[str, str]] = []
+
+    if prompt.system:
+        messages.append({"role": "system", "content": prompt.system})
+
+    if conversation is not None:
+        for prev in conversation.responses:
+            messages.append({"role": "user", "content": prev.prompt.prompt or ""})
+            messages.append(
+                {"role": "assistant", "content": prev.text() or ""},  # type: ignore[union-attr]
+            )
+
+    messages.append({"role": "user", "content": prompt.prompt or ""})
+    return messages

llm_coreml-0.1.0/src/llm_coreml/registry.py

@@ -0,0 +1,57 @@
+"""JSON registry for CoreML model configurations."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+import llm
+
+
+def get_registry_path() -> Path:
+    """Return the path to the registry JSON file."""
+    return Path(llm.user_dir()) / "llm-coreml.json"
+
+
+def _read_registry() -> dict[str, dict[str, str]]:
+    path = get_registry_path()
+    if not path.exists():
+        return {}
+    return json.loads(path.read_text())  # type: ignore[no-any-return]
+
+
+def _write_registry(data: dict[str, Any]) -> None:
+    path = get_registry_path()
+    path.write_text(json.dumps(data, indent=2) + "\n")
+
+
+def add_model(name: str, path: str, tokenizer: str, compute_units: str = "all") -> None:
+    """Register a model with the given name, path, tokenizer ID, and compute units."""
+    registry = _read_registry()
+    registry[name] = {
+        "path": path,
+        "tokenizer": tokenizer,
+        "compute_units": compute_units,
+    }
+    _write_registry(registry)
+
+
+def remove_model(name: str) -> bool:
+    """Remove a model by name. Returns True if it existed."""
+    registry = _read_registry()
+    if name not in registry:
+        return False
+    del registry[name]
+    _write_registry(registry)
+    return True
+
+
+def list_models() -> dict[str, dict[str, str]]:
+    """Return all registered models."""
+    return _read_registry()
+
+
+def get_model(name: str) -> dict[str, str] | None:
+    """Return config for a single model, or None if not found."""
+    return _read_registry().get(name)