add multimodal embedder support for vllm and mlx
This commit is contained in:
Yiorgis Gozadinos
parent
861c36942b
commit
37215f1a00
8 changed files with 606 additions and 8 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
147
haiku_rag_slim/haiku/rag/embeddings/mlx.py
Normal file
147
haiku_rag_slim/haiku/rag/embeddings/mlx.py
Normal file
|
|
@ -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}")
|
||||
119
haiku_rag_slim/haiku/rag/embeddings/vllm.py
Normal file
119
haiku_rag_slim/haiku/rag/embeddings/vllm.py
Normal file
|
|
@ -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}")
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
106
uv.lock
106
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"
|
||||
|
|
|
|||
Loading…
Reference in a new issue