diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ab06cc8..c7f8955e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,10 @@ ### Added +- **Multimodal embedder support (`provider="mlx"` and `provider="vllm"`).** `EmbedderWrapper` gains `supports_images: bool`, `embed_image_query`, and `embed_images`. Two pluggable paths: + - `provider="mlx"` — Apple Silicon, in-process via the new `[mlx]` optional extra (env-marker-guarded so `uv sync --all-extras` works on Linux/Windows/Intel-Mac without resolver errors). Loads any HF repo that ships an MLX `load_model.py` (default tested model: `jinaai/jina-embeddings-v4-mlx-8bit`, 2048-dim). + - `provider="vllm"` — cross-platform, talks HTTP to a vLLM server's OpenAI-compatible `/v1/embeddings` endpoint with vLLM's `messages` superset (text or `image_url` content parts, base64 data URIs). Works with `Qwen/Qwen3-VL-Embedding-8B` and `jinaai/jina-embeddings-v4`. No Python ML deps added — uses `httpx`. + - Text-only providers (`ollama`, `openai`, `cohere`, `sentence-transformers`) report `supports_images=False` and raise a clear error if image methods are called. - **Picture-handling mode enum.** `processing.pictures: "none" | "description" | "image"` (default `"none"`) replaces the legacy `generate_picture_images` + `picture_description.enabled` pair. - `none` — no picture bytes, no VLM, just structural picture rows. - `description` — VLM runs at ingest, descriptions woven into chunk text, bytes also retained on `document_items.picture_data` so vision QA can be turned on later without reingesting. diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index 75155484..9b91989b 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from pydantic_ai.embeddings import Embedder from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel @@ -8,18 +8,34 @@ from pydantic_ai.providers.openai import OpenAIProvider from haiku.rag.config import AppConfig, Config if TYPE_CHECKING: + from PIL import Image as PILImage + from haiku.rag.store.models.chunk import Chunk -class EmbedderWrapper: - """Wrapper around pydantic-ai Embedder with explicit query/document methods.""" +ImageInput = "bytes | PILImage.Image" - def __init__(self, embedder: Embedder, vector_dim: int): + +class EmbedderWrapper: + """Wrapper around pydantic-ai Embedder with explicit query/document methods. + + Subclasses set ``supports_images = True`` and override the image methods + when the underlying model can encode pictures into the same vector space. + """ + + supports_images: bool = False + + def __init__(self, embedder: Embedder | None, vector_dim: int): self._embedder = embedder self._vector_dim = vector_dim + @property + def vector_dim(self) -> int: + return self._vector_dim + async def embed_query(self, text: str) -> list[float]: """Embed a search query.""" + assert self._embedder is not None result = await self._embedder.embed_query(text) return list(result.embeddings[0]) @@ -27,9 +43,26 @@ class EmbedderWrapper: """Embed documents/chunks for indexing.""" if not texts: return [] + assert self._embedder is not None result = await self._embedder.embed_documents(texts) return [list(e) for e in result.embeddings] + async def embed_image_query(self, image: "Any") -> list[float]: + """Embed a single image as a search query.""" + raise NotImplementedError( + f"{type(self).__name__} does not support image embedding. " + "Configure a multimodal provider (e.g. provider='mlx' or " + "provider='vllm')." + ) + + async def embed_images(self, images: list["Any"]) -> list[list[float]]: + """Batch-embed images for indexing into the same vector space as text.""" + raise NotImplementedError( + f"{type(self).__name__} does not support image embedding. " + "Configure a multimodal provider (e.g. provider='mlx' or " + "provider='vllm')." + ) + def contextualize(chunks: list["Chunk"]) -> list[str]: """Prepare chunk content for embedding/FTS by adding context. @@ -148,4 +181,15 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper: Embedder(f"sentence-transformers:{model_name}"), vector_dim ) + if provider == "mlx": + from haiku.rag.embeddings.mlx import MLXEmbedder + + return MLXEmbedder(model_name, vector_dim) + + if provider == "vllm": + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + base_url = embedding_model.base_url or "http://localhost:8000/v1" + return VLLMMultimodalEmbedder(model_name, vector_dim, base_url=base_url) + raise ValueError(f"Unsupported embedding provider: {provider}") diff --git a/haiku_rag_slim/haiku/rag/embeddings/mlx.py b/haiku_rag_slim/haiku/rag/embeddings/mlx.py new file mode 100644 index 00000000..26ad8229 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/embeddings/mlx.py @@ -0,0 +1,147 @@ +"""In-process multimodal embedder backed by Apple's MLX framework. + +Loads a Hugging Face repo that ships an MLX-formatted weight set plus a +``load_model.py`` helper (e.g. ``jinaai/jina-embeddings-v4-mlx-8bit``). +Apple Silicon only — the underlying ``mlx`` / ``mlx-lm`` packages don't +have wheels on other platforms and are guarded by environment markers in +the ``[mlx]`` extra. +""" + +import asyncio +import io +import platform +import sys +from typing import TYPE_CHECKING, Any + +from haiku.rag.embeddings import EmbedderWrapper + +if TYPE_CHECKING: + from PIL import Image as PILImage + + +_DEFAULT_TEXT_PROMPT = "<|im_start|>user\n{text}<|im_end|>" +_DEFAULT_IMAGE_PROMPT = ( + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + "Describe the image.<|im_end|>" +) +_DEFAULT_PROCESSOR_REPO = "jinaai/jina-embeddings-v4" + + +class MLXEmbedder(EmbedderWrapper): + supports_images = True + + def __init__( + self, + model_name: str, + vector_dim: int, + processor_repo: str | None = None, + ): + if sys.platform != "darwin" or platform.machine() != "arm64": + raise RuntimeError( + "provider='mlx' requires Apple Silicon (macOS arm64). " + "On other platforms use provider='vllm' against a vLLM server." + ) + super().__init__(embedder=None, vector_dim=vector_dim) + self._model_name = model_name + self._processor_repo = processor_repo or _DEFAULT_PROCESSOR_REPO + self._model: Any | None = None + self._processor: Any | None = None + + def _ensure_loaded(self) -> tuple[Any, Any]: + if self._model is not None and self._processor is not None: + return self._model, self._processor + + from huggingface_hub import snapshot_download + from transformers import AutoProcessor + + model_dir = snapshot_download(self._model_name) + if model_dir not in sys.path: + sys.path.insert(0, model_dir) + from load_model import ( # type: ignore[import-not-found] # ty: ignore[unresolved-import] + load_mlx_model, + ) + + self._model = load_mlx_model(model_dir) + self._processor = AutoProcessor.from_pretrained( + self._processor_repo, trust_remote_code=True + ) + return self._model, self._processor + + async def embed_query(self, text: str) -> list[float]: + embeddings = await self.embed_documents([text]) + return embeddings[0] + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + if not texts: + return [] + return await asyncio.to_thread(self._encode_texts, texts) + + async def embed_image_query(self, image: "bytes | PILImage.Image") -> list[float]: + embeddings = await self.embed_images([image]) + return embeddings[0] + + async def embed_images( + self, images: list["bytes | PILImage.Image"] + ) -> list[list[float]]: + if not images: + return [] + return await asyncio.to_thread(self._encode_images, images) + + def _encode_texts(self, texts: list[str]) -> list[list[float]]: + import mlx.core as mx # ty: ignore[unresolved-import,unused-ignore-comment] + + model, processor = self._ensure_loaded() + prompts = [_DEFAULT_TEXT_PROMPT.format(text=t) for t in texts] + inputs = processor( + text=prompts, + return_tensors="np", + padding=True, + truncation=True, + max_length=512, + ) + embeddings = model.encode_text( + input_ids=mx.array(inputs["input_ids"]), + attention_mask=mx.array(inputs["attention_mask"]), + task="retrieval", + ) + mx.eval(embeddings) + return [list(map(float, row)) for row in embeddings] + + def _encode_images( + self, images: list["bytes | PILImage.Image"] + ) -> list[list[float]]: + import mlx.core as mx # ty: ignore[unresolved-import,unused-ignore-comment] + from PIL import Image as PILImageModule + + model, processor = self._ensure_loaded() + pil_images = [_to_pil(img) for img in images] + out: list[list[float]] = [] + for pil_image in pil_images: + assert isinstance(pil_image, PILImageModule.Image) + inputs = processor( + text=[_DEFAULT_IMAGE_PROMPT], + images=[pil_image], + return_tensors="np", + padding=True, + ) + pixel_values = inputs["pixel_values"] + embedding = model.encode_image( + input_ids=mx.array(inputs["input_ids"]), + pixel_values=mx.array(pixel_values.reshape(-1, pixel_values.shape[-1])), + image_grid_thw=[tuple(r) for r in inputs["image_grid_thw"]], + attention_mask=mx.array(inputs["attention_mask"]), + task="retrieval", + ) + mx.eval(embedding) + out.append([float(x) for x in embedding[0]]) + return out + + +def _to_pil(image: "bytes | PILImage.Image") -> "PILImage.Image": + from PIL import Image as PILImageModule + + if isinstance(image, bytes): + return PILImageModule.open(io.BytesIO(image)).convert("RGB") + if isinstance(image, PILImageModule.Image): + return image.convert("RGB") if image.mode != "RGB" else image + raise TypeError(f"Unsupported image type: {type(image)!r}") diff --git a/haiku_rag_slim/haiku/rag/embeddings/vllm.py b/haiku_rag_slim/haiku/rag/embeddings/vllm.py new file mode 100644 index 00000000..0f7f1871 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/embeddings/vllm.py @@ -0,0 +1,119 @@ +"""Multimodal embedder backed by a vLLM OpenAI-compatible HTTP server. + +vLLM's ``/v1/embeddings`` endpoint is a superset of OpenAI's: when the +request body uses a ``messages`` array (instead of ``input``), the server +treats it as a chat-style multimodal embedding request and accepts +``image_url`` content parts carrying base64 data URIs. Models like +``Qwen/Qwen3-VL-Embedding-8B`` and ``jinaai/jina-embeddings-v4`` ship with +chat templates that map both text and image inputs into a shared vector +space. +""" + +import asyncio +import base64 +import io +from typing import TYPE_CHECKING, Any + +import httpx + +from haiku.rag.embeddings import EmbedderWrapper + +if TYPE_CHECKING: + from PIL import Image as PILImage + + +class VLLMMultimodalEmbedder(EmbedderWrapper): + supports_images = True + + def __init__( + self, + model_name: str, + vector_dim: int, + base_url: str, + api_key: str | None = None, + timeout: float = 60.0, + ): + super().__init__(embedder=None, vector_dim=vector_dim) + self._model_name = model_name + self._base_url = base_url.rstrip("/") + self._api_key = api_key + self._timeout = timeout + + def _headers(self) -> dict[str, str]: + headers = {"Content-Type": "application/json"} + if self._api_key: + headers["Authorization"] = f"Bearer {self._api_key}" + return headers + + async def _post_messages(self, content: list[dict[str, Any]]) -> list[float]: + body = { + "model": self._model_name, + "messages": [{"role": "user", "content": content}], + "encoding_format": "float", + } + try: + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.post( + f"{self._base_url}/embeddings", + json=body, + headers=self._headers(), + ) + response.raise_for_status() + payload = response.json() + except httpx.ConnectError as e: + raise ValueError( + f"Could not connect to vLLM at {self._base_url}. " + f"Ensure the service is running. Error: {e}" + ) from e + except httpx.TimeoutException as e: + raise ValueError( + f"Request to vLLM timed out after {self._timeout}s. Error: {e}" + ) from e + except httpx.HTTPStatusError as e: + if e.response.status_code == 401: + raise ValueError( + "Authentication failed against vLLM. Check the API key." + ) from e + raise ValueError(f"HTTP error from vLLM: {e}") from e + + data = payload.get("data") or [] + if not data: + raise ValueError(f"vLLM returned no embeddings: {payload}") + return list(data[0]["embedding"]) + + async def embed_query(self, text: str) -> list[float]: + return await self._post_messages([{"type": "text", "text": text}]) + + async def embed_documents(self, texts: list[str]) -> list[list[float]]: + if not texts: + return [] + return await asyncio.gather(*(self.embed_query(t) for t in texts)) + + async def embed_image_query(self, image: "bytes | PILImage.Image") -> list[float]: + return await self._post_messages( + [{"type": "image_url", "image_url": {"url": _to_data_uri(image)}}] + ) + + async def embed_images( + self, images: list["bytes | PILImage.Image"] + ) -> list[list[float]]: + if not images: + return [] + return await asyncio.gather(*(self.embed_image_query(img) for img in images)) + + +def _to_data_uri(image: "bytes | PILImage.Image") -> str: + """Render an image as a ``data:image/png;base64,...`` URI.""" + if isinstance(image, bytes): + return f"data:image/png;base64,{base64.b64encode(image).decode('ascii')}" + + from PIL import Image as PILImageModule + + if isinstance(image, PILImageModule.Image): + buf = io.BytesIO() + image.save(buf, format="PNG") + return ( + f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}" + ) + + raise TypeError(f"Unsupported image type: {type(image)!r}") diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 70b71948..44495185 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -45,6 +45,14 @@ dependencies = [ docling = ["docling>=2.84.0", "opencv-python-headless>=4.13.0.92"] # Embedding providers voyageai = ["pydantic-ai-slim[voyageai]"] +mlx = [ + "mlx>=0.20.0; sys_platform == 'darwin' and platform_machine == 'arm64'", + "mlx-lm>=0.20.0; sys_platform == 'darwin' and platform_machine == 'arm64'", + "huggingface-hub>=0.27.0", + "pillow>=10.0.0", + "transformers>=4.40.0", + "peft>=0.10.0", +] # Rerankers mxbai = ["mxbai-rerank>=0.1.6"] cohere = ["cohere>=5.21.1"] diff --git a/pyproject.toml b/pyproject.toml index fe94c4c4..ccb1127e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ classifiers = [ ] dependencies = [ - "haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.45.0", + "haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,mlx]==0.45.0", ] [project.scripts] diff --git a/tests/test_embedder.py b/tests/test_embedder.py index ee61557d..da4da4c3 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -1,3 +1,5 @@ +import platform +import sys from pathlib import Path import numpy as np @@ -216,3 +218,177 @@ async def test_embed_chunks_preserves_all_fields(allow_model_requests): assert embedded[0].document_title == "Test Document" assert embedded[0].document_meta == {"author": "Test"} assert embedded[0].embedding is not None + + +# B1: multimodal embedder support + + +def _ollama_text_only_config(): + return AppConfig( + embeddings=EmbeddingsConfig( + model=EmbeddingModelConfig( + provider="ollama", name="mxbai-embed-large", vector_dim=1024 + ) + ) + ) + + +async def test_text_only_embedder_does_not_support_images(): + embedder = get_embedder(_ollama_text_only_config()) + assert embedder.supports_images is False + with pytest.raises(NotImplementedError, match="multimodal provider"): + await embedder.embed_image_query(b"\x89PNG\r\n\x1a\n") + with pytest.raises(NotImplementedError, match="multimodal provider"): + await embedder.embed_images([b"\x89PNG\r\n\x1a\n"]) + + +async def test_vllm_embed_text_request_shape(monkeypatch): + """vLLM text embedding posts a `messages` array with a single text part.""" + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + captured: dict = {} + + class FakeResponse: + def raise_for_status(self): + pass + + def json(self): + return {"data": [{"embedding": [0.1, 0.2, 0.3]}]} + + class FakeAsyncClient: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + async def post(self, url, json, headers): + captured["url"] = url + captured["body"] = json + return FakeResponse() + + monkeypatch.setattr("httpx.AsyncClient", FakeAsyncClient) + + embedder = VLLMMultimodalEmbedder( + model_name="Qwen/Qwen3-VL-Embedding-2B", + vector_dim=2048, + base_url="http://localhost:8000/v1", + ) + vec = await embedder.embed_query("a photo of a cat") + assert vec == [0.1, 0.2, 0.3] + assert captured["url"] == "http://localhost:8000/v1/embeddings" + body = captured["body"] + assert body["model"] == "Qwen/Qwen3-VL-Embedding-2B" + assert body["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "a photo of a cat"}]} + ] + assert body["encoding_format"] == "float" + + +async def test_vllm_embed_image_request_shape(monkeypatch): + """vLLM image embedding posts an `image_url` content part with a data: URI.""" + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + captured: dict = {} + + class FakeResponse: + def raise_for_status(self): + pass + + def json(self): + return {"data": [{"embedding": [0.4] * 4}]} + + class FakeAsyncClient: + def __init__(self, *args, **kwargs): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + async def post(self, url, json, headers): + captured["body"] = json + return FakeResponse() + + monkeypatch.setattr("httpx.AsyncClient", FakeAsyncClient) + + embedder = VLLMMultimodalEmbedder( + model_name="some-model", + vector_dim=4, + base_url="http://localhost:8000/v1", + ) + raw = b"\x89PNG\r\n\x1a\nfake" + vec = await embedder.embed_image_query(raw) + assert vec == [0.4, 0.4, 0.4, 0.4] + content = captured["body"]["messages"][0]["content"] + assert len(content) == 1 + assert content[0]["type"] == "image_url" + url = content[0]["image_url"]["url"] + assert url.startswith("data:image/png;base64,") + + +async def test_vllm_supports_images_flag(): + from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder + + embedder = VLLMMultimodalEmbedder( + model_name="x", vector_dim=2, base_url="http://localhost:8000/v1" + ) + assert embedder.supports_images is True + + +async def test_vllm_get_embedder_routes_to_multimodal(): + config = AppConfig( + embeddings=EmbeddingsConfig( + model=EmbeddingModelConfig( + provider="vllm", + name="Qwen/Qwen3-VL-Embedding-2B", + vector_dim=2048, + base_url="http://my-vllm:8000/v1", + ) + ) + ) + embedder = get_embedder(config) + assert embedder.supports_images is True + assert embedder._base_url == "http://my-vllm:8000/v1" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + +def test_mlx_platform_guard_raises_off_apple_silicon(monkeypatch): + monkeypatch.setattr("sys.platform", "linux") + from haiku.rag.embeddings.mlx import MLXEmbedder + + with pytest.raises(RuntimeError, match="Apple Silicon"): + MLXEmbedder(model_name="any", vector_dim=2048) + + +def test_mlx_smoke(): + """End-to-end smoke against a real MLX model. Requires the [mlx] extra.""" + pytest.importorskip("mlx") + if not (sys.platform == "darwin" and platform.machine() == "arm64"): + pytest.skip("MLX path requires Apple Silicon") + + from PIL import Image as PILImageModule + + from haiku.rag.embeddings.mlx import MLXEmbedder + + embedder = MLXEmbedder( + model_name="jinaai/jina-embeddings-v4-mlx-8bit", + vector_dim=2048, + ) + + async def run(): + text_vec = await embedder.embed_query("a photo of a cat") + assert len(text_vec) == 2048 + img_vec = await embedder.embed_image_query( + PILImageModule.new("RGB", (224, 224), "red") + ) + assert len(img_vec) == 2048 + assert sum(a * b for a, b in zip(text_vec, img_vec)) != 0.0 + + import asyncio + + asyncio.run(run()) diff --git a/uv.lock b/uv.lock index 7f8f2d33..d340925c 100644 --- a/uv.lock +++ b/uv.lock @@ -1421,7 +1421,7 @@ name = "haiku-rag" version = "0.45.0" source = { editable = "." } dependencies = [ - { name = "haiku-rag-slim", extra = ["cohere", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] }, + { name = "haiku-rag-slim", extra = ["cohere", "docling", "mlx", "mxbai", "tui", "voyageai", "zeroentropy"] }, ] [package.optional-dependencies] @@ -1448,7 +1448,7 @@ dev = [ [package.metadata] requires-dist = [ - { name = "haiku-rag-slim", extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "tui"], editable = "haiku_rag_slim" }, + { name = "haiku-rag-slim", extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "tui", "mlx"], editable = "haiku_rag_slim" }, { name = "textual", marker = "extra == 'tui'", specifier = ">=1.0.0" }, ] provides-extras = ["tui"] @@ -1548,6 +1548,14 @@ jina = [ mistral = [ { name = "pydantic-ai-slim", extra = ["mistral"] }, ] +mlx = [ + { name = "huggingface-hub" }, + { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, + { name = "mlx-lm", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, + { name = "peft" }, + { name = "pillow" }, + { name = "transformers" }, +] mxbai = [ { name = "mxbai-rerank" }, ] @@ -1572,12 +1580,17 @@ requires-dist = [ { name = "docling-core", specifier = ">=2.71.0,<2.72" }, { name = "haiku-skills", specifier = ">=0.16.0" }, { name = "httpx", specifier = ">=0.28.1" }, + { name = "huggingface-hub", marker = "extra == 'mlx'", specifier = ">=0.27.0" }, { name = "jinja2", specifier = ">=3.1.0" }, { name = "jsonpatch", specifier = ">=1.33" }, { name = "lancedb", specifier = "==0.30.2" }, + { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'mlx'", specifier = ">=0.20.0" }, + { name = "mlx-lm", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'mlx'", specifier = ">=0.20.0" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" }, { name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.13.0.92" }, { name = "pathspec", specifier = ">=1.0.4" }, + { name = "peft", marker = "extra == 'mlx'", specifier = ">=0.10.0" }, + { name = "pillow", marker = "extra == 'mlx'", specifier = ">=10.0.0" }, { name = "pydantic", specifier = ">=2.12.5" }, { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'" }, { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'" }, @@ -1595,12 +1608,13 @@ requires-dist = [ { name = "textual-image", marker = "extra == 'tui'", specifier = ">=0.8.5" }, { name = "torch", marker = "extra == 'jina'", specifier = ">=2.0.0" }, { name = "transformers", marker = "extra == 'jina'", specifier = ">=4.40.0" }, + { name = "transformers", marker = "extra == 'mlx'", specifier = ">=4.40.0" }, { name = "typer", specifier = ">=0.21.0,<0.22.0" }, { name = "watchfiles", specifier = ">=1.1.1" }, { name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a11" }, { name = "zstandard", marker = "python_full_version < '3.14'", specifier = ">=0.23.0" }, ] -provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"] +provides-extras = ["docling", "voyageai", "mlx", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"] [[package]] name = "haiku-skills" @@ -2460,6 +2474,53 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/54/662a4743aa81d9582ee9339d4ffa3c8fd40a4965e033d77b9da9774d3960/mkdocs_material_extensions-1.3.1-py3-none-any.whl", hash = "sha256:adff8b62700b25cb77b53358dad940f3ef973dd6db797907c49e3c2ef3ab4e31", size = 8728, upload-time = "2023-11-22T19:09:43.465Z" }, ] +[[package]] +name = "mlx" +version = "0.31.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/47/5f33906cb03d6a378a697cd2d2641a26b37dea17ee3d9124d7e39e8eca01/mlx-0.31.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:e5067aaf2be1f3d7bba5be52348775804f111173c1ed04639618fd713b1a530f", size = 584863, upload-time = "2026-04-22T03:14:38.211Z" }, + { url = "https://files.pythonhosted.org/packages/08/e7/a851a451b1327af9fb4df3991b9ae87d066b6f6630e854af55c288b0995a/mlx-0.31.2-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:edb9797db7d852477ca1c99708058654ee860d4148fe5765f0d55528e2b1aa22", size = 584860, upload-time = "2026-04-22T03:14:39.746Z" }, + { url = "https://files.pythonhosted.org/packages/3b/15/0d1dc0597644e5e7b011ca954ba0c47e13cd880a3b909b0c3f1b4d8bf8f1/mlx-0.31.2-cp312-cp312-macosx_26_0_arm64.whl", hash = "sha256:51ca102db641b01e7cb083ce8ecb580e281530a141a7ca12544bb370641630ae", size = 584887, upload-time = "2026-04-22T03:14:41.585Z" }, + { url = "https://files.pythonhosted.org/packages/a3/3f/888f8664d4f8e23a1363a5f50024be5216e199ab7ad0ba20988c7ed6d729/mlx-0.31.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:1b3fb0dda955b0d552ce57bdd6f42b3309ab21b067e40587d6848443d307e91f", size = 584796, upload-time = "2026-04-22T03:14:47.215Z" }, + { url = "https://files.pythonhosted.org/packages/dd/14/e9cd18b51f9e1dbcb060eec0fafc2d2428c8e1eacd9b0a02d7c5ce75b661/mlx-0.31.2-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:34b0171cd9eb5c43fdd82091f6135d6ccc5a065363a4a3e68fac64fb4e53d37c", size = 584790, upload-time = "2026-04-22T03:14:48.519Z" }, + { url = "https://files.pythonhosted.org/packages/ca/20/c6c5fb998c7834d094b2bfb9f003b5246cb270f0266da055c55546c34999/mlx-0.31.2-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:c05981684279a8935d58b0dde3ea5b02d210c3bad3319aa0e9934ec2df165752", size = 584795, upload-time = "2026-04-22T03:14:49.904Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f5/e63f6a9316ded2d14a8ebc7a9ca25734c784e8c54d064a78b4dceeacec0e/mlx-0.31.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:a13c9ce23c3deef6aa5a09315e7953e1a5dc311e851fa16fc74c81fb2509c0b9", size = 588417, upload-time = "2026-04-22T03:14:54.094Z" }, + { url = "https://files.pythonhosted.org/packages/31/50/9d0c03ea3134cd85c132df7b0e4b75e6344bd8b4881a0b9c465cfa27f724/mlx-0.31.2-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:b0764bf11fc3a71dee988e19275eef67775cab63112d8bb7ef173ca8b2a1247c", size = 588421, upload-time = "2026-04-22T03:14:55.898Z" }, + { url = "https://files.pythonhosted.org/packages/ef/5b/d364cc793bcb504621313acb55627cf0d5403ab2e0a594aa081cdbe4591f/mlx-0.31.2-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:59ccbd0f0044d4f97f11ebcbf0c480bc9e962935fd96275f120954afea65be8a", size = 588384, upload-time = "2026-04-22T03:14:57.439Z" }, +] + +[[package]] +name = "mlx-lm" +version = "0.29.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jinja2", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "mlx", marker = "sys_platform == 'darwin'" }, + { name = "numpy", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "protobuf", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "pyyaml", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "sentencepiece", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "transformers", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e3/62/f46e1355256a114808517947f8e83ad6be310c7288c551db0fa678f47923/mlx_lm-0.29.1.tar.gz", hash = "sha256:b99180d8f33d33a077b814e550bfb2d8a59ae003d668fd1f4b3fff62a381d34b", size = 232302, upload-time = "2025-12-16T16:58:27.959Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/53/913099c91d384e115ea078325efd9a0bc1ea3eb3458c694b4596cbd267f2/mlx_lm-0.29.1-py3-none-any.whl", hash = "sha256:440941b3054c2a2216e97615de584cc90fa1ea874782e20699b9895721fad8dc", size = 324884, upload-time = "2025-12-16T16:58:26.36Z" }, +] + +[[package]] +name = "mlx-metal" +version = "0.31.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/69/fe3b783ebe999f3118234e1e940feb622518bfb1dea6ac5d13b1d36a8449/mlx_metal-0.31.2-py3-none-macosx_14_0_arm64.whl", hash = "sha256:b25385bcee18fc194092255b8b53b9a3d8489eb650e59160f1b57aadd07aa2dc", size = 40055588, upload-time = "2026-04-22T03:14:14.43Z" }, + { url = "https://files.pythonhosted.org/packages/4f/5d/4c690d5b93c30ba002656c37363159d978705bf8eb801b8481840fb942c2/mlx_metal-0.31.2-py3-none-macosx_15_0_arm64.whl", hash = "sha256:e9d4e5fce6ca10a87a0e388597f99519ad594d09e674708b5312bd8bd4f5997d", size = 40053220, upload-time = "2026-04-22T03:14:18.048Z" }, + { url = "https://files.pythonhosted.org/packages/99/82/11fd62a8d7a3e96e5c43220b17de0151e3f10101f8bb3b865f5bd9cdd074/mlx_metal-0.31.2-py3-none-macosx_26_0_arm64.whl", hash = "sha256:84ffb60ee503f03eb684f5fb168d5cff31e2a16b7f27c1731eaf7662bd6e9b46", size = 55792151, upload-time = "2026-04-22T03:14:22.059Z" }, +] + [[package]] name = "more-itertools" version = "10.8.0" @@ -3227,6 +3288,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ef/3c/2c197d226f9ea224a9ab8d197933f9da0ae0aac5b6e0f884e2b8d9c8e9f7/pathspec-1.0.4-py3-none-any.whl", hash = "sha256:fb6ae2fd4e7c921a165808a552060e722767cfa526f99ca5156ed2ce45a5c723", size = 55206, upload-time = "2026-01-27T03:59:45.137Z" }, ] +[[package]] +name = "peft" +version = "0.19.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "accelerate" }, + { name = "huggingface-hub" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "psutil" }, + { name = "pyyaml" }, + { name = "safetensors" }, + { name = "torch" }, + { name = "tqdm" }, + { name = "transformers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/86/cf/037f1e3d5186496c05513a6754639e2dab3038a05f384284d49a9bd06a2d/peft-0.19.1.tar.gz", hash = "sha256:0d97542fe96dcdaa20d3b81c06f26f988618f416a73544ab23c3618ccb674a40", size = 763738, upload-time = "2026-04-16T15:46:45.105Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/b6/f54d676ed93cc2dd2234c3b172ea9c8c3d7d29361e66b1b23dec57a67465/peft-0.19.1-py3-none-any.whl", hash = "sha256:2113f72a81621b5913ef28f9022204c742df111890c5f49d812716a4a301e356", size = 680692, upload-time = "2026-04-16T15:46:42.886Z" }, +] + [[package]] name = "pillow" version = "12.1.1" @@ -4652,6 +4734,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f8/95/12d226ee4d207cb1f77a216baa7e1a8bae2639733c140abe8d0316d23a18/semchunk-3.2.5-py3-none-any.whl", hash = "sha256:fd09cc5f380bd010b8ca773bd81893f7eaf11d37dd8362a83d46cedaf5dae076", size = 13048, upload-time = "2025-10-28T02:12:36.724Z" }, ] +[[package]] +name = "sentencepiece" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" }, + { url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" }, + { url = "https://files.pythonhosted.org/packages/ba/4a/85fbe1706d4d04a7e826b53f327c4b80f849cf1c7b7c5e31a20a97d8f28b/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dcd8161eee7b41aae57ded06272905dbd680a0a04b91edd0f64790c796b2f706", size = 1943150, upload-time = "2025-08-12T06:59:53.588Z" }, + { url = "https://files.pythonhosted.org/packages/8d/de/5a007fb53b1ab0aafc69d11a5a3dd72a289d5a3e78dcf2c3a3d9b14ffe93/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:097f3394e99456e9e4efba1737c3749d7e23563dd1588ce71a3d007f25475fff", size = 1253641, upload-time = "2025-08-12T06:59:56.562Z" }, + { url = "https://files.pythonhosted.org/packages/4a/b6/08fe2ce819e02ccb0296f4843e3f195764ce9829cbda61b7513f29b95718/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8dd4b477a7b069648d19363aad0cab9bad2f4e83b2d179be668efa672500dc94", size = 1946052, upload-time = "2025-08-12T07:00:08.136Z" }, + { url = "https://files.pythonhosted.org/packages/99/7e/1fb26e8a21613f6200e1ab88824d5d203714162cf2883248b517deb500b7/sentencepiece-0.2.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad8493bea8432dae8d6830365352350f3b4144415a1d09c4c8cb8d30cf3b6c3c", size = 1254857, upload-time = "2025-08-12T07:00:11.021Z" }, + { url = "https://files.pythonhosted.org/packages/24/9c/89eb8b2052f720a612478baf11c8227dcf1dc28cd4ea4c0c19506b5af2a2/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:5d0350b686c320068702116276cfb26c066dc7e65cfef173980b11bb4d606719", size = 1943147, upload-time = "2025-08-12T07:00:21.809Z" }, + { url = "https://files.pythonhosted.org/packages/ea/99/bbe054ebb5a5039457c590e0a4156ed073fb0fe9ce4f7523404dd5b37463/sentencepiece-0.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c83b85ab2d6576607f31df77ff86f28182be4a8de6d175d2c33ca609925f5da1", size = 1253670, upload-time = "2025-08-12T07:00:24.69Z" }, + { url = "https://files.pythonhosted.org/packages/a1/11/5b414b9fae6255b5fb1e22e2ed3dc3a72d3a694e5703910e640ac78346bb/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a19adcec27c524cb7069a1c741060add95f942d1cbf7ad0d104dffa0a7d28a2b", size = 1946081, upload-time = "2025-08-12T07:00:36.97Z" }, + { url = "https://files.pythonhosted.org/packages/03/b0/811dae8fb9f2784e138785d481469788f2e0d0c109c5737372454415f55f/sentencepiece-0.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:477c81505db072b3ab627e7eab972ea1025331bd3a92bacbf798df2b75ea86ec", size = 1254846, upload-time = "2025-08-12T07:00:40.611Z" }, +] + [[package]] name = "setuptools" version = "81.0.0"