Remove the mxbai reranking provider
This commit is contained in:
parent
5e2a928013
commit
144900d385
20 changed files with 35 additions and 192 deletions
|
|
@ -3,7 +3,11 @@
|
||||||
|
|
||||||
### Changed
|
### Changed
|
||||||
|
|
||||||
- `mxbai` extra allows `transformers` 5.x (`<6.0.0`); `MxBAIReranker` supplies the `tokenizer.prepare_for_model` variant `mxbai-rerank` needs when the tokenizer lacks it.
|
- Unknown `reranking.model.provider` raises `ValueError` instead of silently disabling reranking.
|
||||||
|
|
||||||
|
### Removed
|
||||||
|
|
||||||
|
- `mxbai` reranking provider and extra; the `transformers<5.0.0` cap goes with it. Migrate to `provider: cross-encoder` with the same model name (`mixedbread-ai/mxbai-rerank-base-v2`), installed via the `cross-encoder` extra.
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -13,7 +13,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
|
||||||
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
|
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
|
||||||
- **Question answering** — RAG skill with citations (page numbers, section headings)
|
- **Question answering** — RAG skill with citations (page numbers, section headings)
|
||||||
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text
|
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text
|
||||||
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
|
- **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM
|
||||||
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||||
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
|
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
|
||||||
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
|
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,6 @@
|
||||||
# haiku.rag Docker Image
|
# haiku.rag Docker Image
|
||||||
|
|
||||||
The full haiku.rag Docker image includes all features and extras (docling, voyageai, mxbai). You can build it locally using the provided Dockerfile.
|
The full haiku.rag Docker image includes all features and extras (docling, voyageai, cross-encoder). You can build it locally using the provided Dockerfile.
|
||||||
|
|
||||||
## Building the Image
|
## Building the Image
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -322,7 +322,7 @@ It also probes the external endpoints the config uses and reports them under a P
|
||||||
- docling-serve is reachable when used as the converter or chunker (`{base_url}/health`)
|
- docling-serve is reachable when used as the converter or chunker (`{base_url}/health`)
|
||||||
- custom OpenAI-compatible and vLLM endpoints respond (`{base_url}/models`)
|
- custom OpenAI-compatible and vLLM endpoints respond (`{base_url}/models`)
|
||||||
|
|
||||||
SaaS providers (OpenAI, Anthropic, Cohere, Jina, ZeroEntropy, Voyage) are covered by the API-key check rather than a network probe. In-process local models (sentence-transformers, cross-encoder, mxbai, jina-local) have no endpoint and are reported as such.
|
SaaS providers (OpenAI, Anthropic, Cohere, Jina, ZeroEntropy, Voyage) are covered by the API-key check rather than a network probe. In-process local models (sentence-transformers, cross-encoder, jina-local) have no endpoint and are reported as such.
|
||||||
|
|
||||||
Each failure prints the command that fixes it (`rebuild`, `create-index`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring.
|
Each failure prints the command that fixes it (`rebuild`, `create-index`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -97,7 +97,7 @@ embeddings:
|
||||||
|
|
||||||
reranking:
|
reranking:
|
||||||
model:
|
model:
|
||||||
provider: "" # Empty to disable, or mxbai, cohere, zeroentropy, vllm
|
provider: "" # Empty to disable, or cross-encoder, cohere, zeroentropy, vllm
|
||||||
name: ""
|
name: ""
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
|
|
||||||
|
|
@ -390,23 +390,6 @@ Reranking improves search quality by re-ordering the initial search results usin
|
||||||
|
|
||||||
Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below.
|
Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below.
|
||||||
|
|
||||||
### MixedBread AI
|
|
||||||
|
|
||||||
If you installed `haiku.rag` (full package), MxBAI is already included. If you installed `haiku.rag-slim`, add the mxbai extra:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
uv pip install haiku.rag-slim[mxbai]
|
|
||||||
```
|
|
||||||
|
|
||||||
Then configure:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
reranking:
|
|
||||||
model:
|
|
||||||
provider: mxbai
|
|
||||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
|
||||||
```
|
|
||||||
|
|
||||||
### Cohere
|
### Cohere
|
||||||
|
|
||||||
If you installed `haiku.rag` (full package), Cohere is already included. If you installed `haiku.rag-slim`, add the cohere extra:
|
If you installed `haiku.rag` (full package), Cohere is already included. If you installed `haiku.rag-slim`, add the cohere extra:
|
||||||
|
|
@ -523,7 +506,7 @@ Then configure with any HuggingFace model id:
|
||||||
reranking:
|
reranking:
|
||||||
model:
|
model:
|
||||||
provider: cross-encoder
|
provider: cross-encoder
|
||||||
name: BAAI/bge-reranker-v2-m3
|
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||||
```
|
```
|
||||||
|
|
||||||
Other tested models: `Qwen/Qwen3-Reranker-0.6B`, `cross-encoder/ms-marco-MiniLM-L-6-v2`. Any model exposed as a `sentence_transformers.CrossEncoder` works.
|
Other tested models: `BAAI/bge-reranker-v2-m3`, `Qwen/Qwen3-Reranker-0.6B`, `cross-encoder/ms-marco-MiniLM-L-6-v2`. Any model exposed as a `sentence_transformers.CrossEncoder` works.
|
||||||
|
|
|
||||||
|
|
@ -25,14 +25,14 @@ uv pip install haiku.rag-slim
|
||||||
uv pip install haiku.rag-slim[docling]
|
uv pip install haiku.rag-slim[docling]
|
||||||
|
|
||||||
# With specific providers
|
# With specific providers
|
||||||
uv pip install haiku.rag-slim[docling,voyageai,mxbai]
|
uv pip install haiku.rag-slim[docling,voyageai,cross-encoder]
|
||||||
```
|
```
|
||||||
|
|
||||||
The slim package has minimal dependencies and lets you install only what you need:
|
The slim package has minimal dependencies and lets you install only what you need:
|
||||||
|
|
||||||
- `docling` - PDF, DOCX, PPTX, images, and other document formats
|
- `docling` - PDF, DOCX, PPTX, images, and other document formats
|
||||||
- `voyageai` - VoyageAI embeddings
|
- `voyageai` - VoyageAI embeddings
|
||||||
- `mxbai` - MixedBread AI reranking
|
- `cross-encoder` - Local reranking via sentence-transformers
|
||||||
- `cohere` - Cohere reranking
|
- `cohere` - Cohere reranking
|
||||||
- `zeroentropy` - Zero Entropy reranking
|
- `zeroentropy` - Zero Entropy reranking
|
||||||
- `tui` - Terminal UI for `chat` and `inspect` commands
|
- `tui` - Terminal UI for `chat` and `inspect` commands
|
||||||
|
|
|
||||||
|
|
@ -17,7 +17,7 @@ embeddings:
|
||||||
|
|
||||||
reranking:
|
reranking:
|
||||||
model:
|
model:
|
||||||
provider: mxbai
|
provider: cross-encoder
|
||||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
|
|
||||||
|
|
@ -22,7 +22,7 @@ processing:
|
||||||
|
|
||||||
reranking:
|
reranking:
|
||||||
model:
|
model:
|
||||||
provider: mxbai
|
provider: cross-encoder
|
||||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ embeddings:
|
||||||
|
|
||||||
reranking:
|
reranking:
|
||||||
model:
|
model:
|
||||||
provider: mxbai
|
provider: cross-encoder
|
||||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
|
|
||||||
|
|
@ -63,12 +63,12 @@ class TestBuildExperimentMetadata:
|
||||||
def test_with_reranker(self) -> None:
|
def test_with_reranker(self) -> None:
|
||||||
config = AppConfig()
|
config = AppConfig()
|
||||||
config.reranking.model = ModelConfig(
|
config.reranking.model = ModelConfig(
|
||||||
provider="mxbai", name="mixedbread-ai/mxbai-rerank-base-v2"
|
provider="cross-encoder", name="mixedbread-ai/mxbai-rerank-base-v2"
|
||||||
)
|
)
|
||||||
result = build_experiment_metadata(
|
result = build_experiment_metadata(
|
||||||
dataset_key="test", test_cases=1, config=config
|
dataset_key="test", test_cases=1, config=config
|
||||||
)
|
)
|
||||||
assert result["rerank_provider"] == "mxbai"
|
assert result["rerank_provider"] == "cross-encoder"
|
||||||
assert result["rerank_model"] == "mixedbread-ai/mxbai-rerank-base-v2"
|
assert result["rerank_model"] == "mixedbread-ai/mxbai-rerank-base-v2"
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -35,7 +35,7 @@ Adds support for 40+ file formats including PDF, DOCX, HTML, and more.
|
||||||
- `voyageai` - VoyageAI embeddings
|
- `voyageai` - VoyageAI embeddings
|
||||||
|
|
||||||
**Rerankers:**
|
**Rerankers:**
|
||||||
- `mxbai` - MixedBread AI
|
- `cross-encoder` - Local reranking via sentence-transformers
|
||||||
- `cohere` - Cohere
|
- `cohere` - Cohere
|
||||||
- `zeroentropy` - Zero Entropy
|
- `zeroentropy` - Zero Entropy
|
||||||
|
|
||||||
|
|
@ -51,7 +51,7 @@ Adds support for 40+ file formats including PDF, DOCX, HTML, and more.
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
# Common combinations
|
# Common combinations
|
||||||
uv pip install haiku.rag-slim[docling,anthropic,mxbai]
|
uv pip install haiku.rag-slim[docling,anthropic,cross-encoder]
|
||||||
uv pip install haiku.rag-slim[docling,groq]
|
uv pip install haiku.rag-slim[docling,groq]
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -28,7 +28,7 @@ async def download_models(
|
||||||
- Docling models
|
- Docling models
|
||||||
- HuggingFace tokenizer
|
- HuggingFace tokenizer
|
||||||
- Sentence-transformers embedder (if configured)
|
- Sentence-transformers embedder (if configured)
|
||||||
- HuggingFace reranker models (mxbai, jina-local)
|
- HuggingFace reranker models (cross-encoder, jina-local)
|
||||||
- Ollama models
|
- Ollama models
|
||||||
"""
|
"""
|
||||||
# Docling models
|
# Docling models
|
||||||
|
|
@ -69,14 +69,14 @@ async def download_models(
|
||||||
provider = config.reranking.model.provider
|
provider = config.reranking.model.provider
|
||||||
model_name = config.reranking.model.name
|
model_name = config.reranking.model.name
|
||||||
|
|
||||||
if provider == "mxbai":
|
if provider == "cross-encoder":
|
||||||
try:
|
try:
|
||||||
from mxbai_rerank import MxbaiRerankV2
|
from sentence_transformers import ( # type: ignore[import-not-found]
|
||||||
|
CrossEncoder,
|
||||||
|
)
|
||||||
|
|
||||||
yield DownloadProgress(model=model_name, status="start")
|
yield DownloadProgress(model=model_name, status="start")
|
||||||
await asyncio.to_thread(
|
await asyncio.to_thread(lambda: CrossEncoder(model_name))
|
||||||
MxbaiRerankV2, model_name, disable_transformers_warnings=True
|
|
||||||
)
|
|
||||||
yield DownloadProgress(model=model_name, status="done")
|
yield DownloadProgress(model=model_name, status="done")
|
||||||
except ImportError:
|
except ImportError:
|
||||||
pass
|
pass
|
||||||
|
|
|
||||||
|
|
@ -34,7 +34,7 @@ _PROVIDER_ENV_VARS: dict[str, str] = {
|
||||||
}
|
}
|
||||||
|
|
||||||
# Providers backed by in-process local models — no endpoint to probe.
|
# Providers backed by in-process local models — no endpoint to probe.
|
||||||
_LOCAL_PROVIDERS = {"sentence-transformers", "mxbai", "cross-encoder", "jina-local"}
|
_LOCAL_PROVIDERS = {"sentence-transformers", "cross-encoder", "jina-local"}
|
||||||
|
|
||||||
# Item labels that never yield a standalone chunk: pictures (handled via the
|
# Item labels that never yield a standalone chunk: pictures (handled via the
|
||||||
# image path), headings (folded into chunk context, not embedded alone), and
|
# image path), headings (folded into chunk context, not embedded alone), and
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,3 @@
|
||||||
import os
|
|
||||||
|
|
||||||
from haiku.rag.config import AppConfig, Config
|
from haiku.rag.config import AppConfig, Config
|
||||||
from haiku.rag.reranking.base import RerankerBase
|
from haiku.rag.reranking.base import RerankerBase
|
||||||
|
|
||||||
|
|
@ -12,12 +10,6 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
try:
|
try:
|
||||||
if model.provider == "mxbai":
|
|
||||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
|
||||||
|
|
||||||
os.environ["TOKENIZERS_PARALLELISM"] = "true"
|
|
||||||
return MxBAIReranker()
|
|
||||||
|
|
||||||
if model.provider == "cohere":
|
if model.provider == "cohere":
|
||||||
from haiku.rag.reranking.cohere import CohereReranker
|
from haiku.rag.reranking.cohere import CohereReranker
|
||||||
|
|
||||||
|
|
@ -56,4 +48,4 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
except ImportError: # pragma: no cover
|
except ImportError: # pragma: no cover
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return None
|
raise ValueError(f"Unknown reranking provider: {model.provider}")
|
||||||
|
|
|
||||||
|
|
@ -1,58 +0,0 @@
|
||||||
import asyncio
|
|
||||||
import threading
|
|
||||||
|
|
||||||
import tqdm
|
|
||||||
from mxbai_rerank import MxbaiRerankV2 # pyright: ignore[reportMissingImports]
|
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.reranking.base import RerankerBase
|
|
||||||
from haiku.rag.store.models.chunk import Chunk
|
|
||||||
|
|
||||||
# tqdm's default class lock is a multiprocessing.RLock; constructing it spawns
|
|
||||||
# resource_tracker, which inherits sys.stderr's fileno. Inside Textual's chat
|
|
||||||
# TUI, sys.stderr.fileno() returns -1, landing in fds_to_keep and failing the
|
|
||||||
# fork_exec validation. A threading lock is sufficient since we never share
|
|
||||||
# tqdm progress bars across processes.
|
|
||||||
tqdm.tqdm.set_lock(threading.RLock())
|
|
||||||
|
|
||||||
|
|
||||||
def _prepare_for_model(
|
|
||||||
ids: list[int],
|
|
||||||
pair_ids: list[int] | None = None,
|
|
||||||
max_length: int | None = None,
|
|
||||||
**_,
|
|
||||||
) -> dict[str, list[int]]:
|
|
||||||
# transformers 5.x removed tokenizer.prepare_for_model, which mxbai-rerank
|
|
||||||
# calls with add_special_tokens=False and truncation="only_second"; for that
|
|
||||||
# call pattern it reduces to truncating the pair and concatenating.
|
|
||||||
pair_ids = pair_ids or []
|
|
||||||
if max_length is not None:
|
|
||||||
pair_ids = pair_ids[: max(0, max_length - len(ids))]
|
|
||||||
return {"input_ids": ids + pair_ids}
|
|
||||||
|
|
||||||
|
|
||||||
class MxBAIReranker(RerankerBase):
|
|
||||||
def __init__(self):
|
|
||||||
model_name = (
|
|
||||||
Config.reranking.model.name
|
|
||||||
if Config.reranking.model
|
|
||||||
else "mixedbread-ai/mxbai-rerank-base-v2"
|
|
||||||
)
|
|
||||||
self._client = MxbaiRerankV2(model_name, disable_transformers_warnings=True)
|
|
||||||
if not hasattr(self._client.tokenizer, "prepare_for_model"):
|
|
||||||
self._client.tokenizer.prepare_for_model = _prepare_for_model
|
|
||||||
|
|
||||||
async def _rerank(
|
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
||||||
) -> list[tuple[Chunk, float]]:
|
|
||||||
documents = [chunk.content for chunk in chunks]
|
|
||||||
|
|
||||||
results = await asyncio.to_thread(
|
|
||||||
lambda: self._client.rank(query=query, documents=documents, top_k=top_n)
|
|
||||||
)
|
|
||||||
reranked_chunks = []
|
|
||||||
for result in results:
|
|
||||||
original_chunk = chunks[result.index]
|
|
||||||
reranked_chunks.append((original_chunk, result.score))
|
|
||||||
|
|
||||||
return reranked_chunks
|
|
||||||
|
|
@ -51,7 +51,6 @@ s3 = ["obstore>=0.9,<0.10"]
|
||||||
# Embedding providers
|
# Embedding providers
|
||||||
voyageai = ["pydantic-ai-slim[voyageai]"]
|
voyageai = ["pydantic-ai-slim[voyageai]"]
|
||||||
# Rerankers
|
# Rerankers
|
||||||
mxbai = ["mxbai-rerank>=0.1.6", "transformers>=4.49.0,<6.0.0"]
|
|
||||||
cohere = ["cohere>=5.21.1"]
|
cohere = ["cohere>=5.21.1"]
|
||||||
zeroentropy = ["zeroentropy>=0.1.0a11"]
|
zeroentropy = ["zeroentropy>=0.1.0a11"]
|
||||||
jina = ["transformers>=4.40.0", "torch>=2.0.0"]
|
jina = ["transformers>=4.40.0", "torch>=2.0.0"]
|
||||||
|
|
|
||||||
|
|
@ -30,7 +30,7 @@ classifiers = [
|
||||||
]
|
]
|
||||||
|
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.65.1",
|
"haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.65.1",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
|
|
|
||||||
|
|
@ -11,7 +11,6 @@ from haiku.rag.store.models.chunk import Chunk
|
||||||
# Providers whose constructor loads a model in-process. Factory-routing tests
|
# Providers whose constructor loads a model in-process. Factory-routing tests
|
||||||
# patch the loader so they assert dispatch without paying the model load.
|
# patch the loader so they assert dispatch without paying the model load.
|
||||||
HEAVY_LOADERS = {
|
HEAVY_LOADERS = {
|
||||||
"mxbai": "MxbaiRerankV2",
|
|
||||||
"jina-local": "AutoModel",
|
"jina-local": "AutoModel",
|
||||||
"cross-encoder": "CrossEncoder",
|
"cross-encoder": "CrossEncoder",
|
||||||
}
|
}
|
||||||
|
|
@ -53,40 +52,6 @@ async def test_reranker_base():
|
||||||
await reranker.rerank("query", chunks)
|
await reranker.rerank("query", chunks)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
@pytest.mark.integration
|
|
||||||
async def test_mxbai_reranker():
|
|
||||||
try:
|
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.config.models import ModelConfig
|
|
||||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
|
||||||
|
|
||||||
Config.reranking.model = ModelConfig(
|
|
||||||
provider="mxbai", name="mixedbread-ai/mxbai-rerank-base-v2"
|
|
||||||
)
|
|
||||||
reranker = MxBAIReranker()
|
|
||||||
reranked = await reranker.rerank(
|
|
||||||
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
|
||||||
)
|
|
||||||
assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
|
|
||||||
assert all(isinstance(score, float) for chunk, score in reranked)
|
|
||||||
Config.reranking.model = None
|
|
||||||
|
|
||||||
except ImportError:
|
|
||||||
pytest.skip("MxBAI package not installed")
|
|
||||||
|
|
||||||
|
|
||||||
def test_mxbai_prepare_for_model_shim():
|
|
||||||
pytest.importorskip("mxbai_rerank")
|
|
||||||
from haiku.rag.reranking.mxbai import _prepare_for_model
|
|
||||||
|
|
||||||
assert _prepare_for_model([1, 2], [3, 4]) == {"input_ids": [1, 2, 3, 4]}
|
|
||||||
assert _prepare_for_model([1, 2], [3, 4, 5], max_length=4) == {
|
|
||||||
"input_ids": [1, 2, 3, 4]
|
|
||||||
}
|
|
||||||
assert _prepare_for_model([1, 2], [3, 4], max_length=2) == {"input_ids": [1, 2]}
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@pytest.mark.vcr()
|
@pytest.mark.vcr()
|
||||||
async def test_cohere_reranker():
|
async def test_cohere_reranker():
|
||||||
|
|
@ -133,14 +98,14 @@ class TestGetReranker:
|
||||||
result = get_reranker(config)
|
result = get_reranker(config)
|
||||||
assert result is None
|
assert result is None
|
||||||
|
|
||||||
def test_unknown_provider_returns_none(self):
|
def test_unknown_provider_raises_error(self):
|
||||||
config = AppConfig(
|
config = AppConfig(
|
||||||
reranking=RerankingConfig(
|
reranking=RerankingConfig(
|
||||||
model=ModelConfig(provider="unknown_provider", name="some-model")
|
model=ModelConfig(provider="unknown_provider", name="some-model")
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
result = get_reranker(config)
|
with pytest.raises(ValueError, match="Unknown reranking provider"):
|
||||||
assert result is None
|
get_reranker(config)
|
||||||
|
|
||||||
def test_vllm_provider_without_base_url_raises_error(self):
|
def test_vllm_provider_without_base_url_raises_error(self):
|
||||||
config = AppConfig(
|
config = AppConfig(
|
||||||
|
|
@ -163,15 +128,6 @@ class TestGetReranker:
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"provider, model_name, class_module, class_name, extra_model_kwargs, expected_attrs, env_vars",
|
"provider, model_name, class_module, class_name, extra_model_kwargs, expected_attrs, env_vars",
|
||||||
[
|
[
|
||||||
(
|
|
||||||
"mxbai",
|
|
||||||
"mixedbread-ai/mxbai-rerank-base-v2",
|
|
||||||
"haiku.rag.reranking.mxbai",
|
|
||||||
"MxBAIReranker",
|
|
||||||
{},
|
|
||||||
{},
|
|
||||||
{},
|
|
||||||
),
|
|
||||||
(
|
(
|
||||||
"cohere",
|
"cohere",
|
||||||
"rerank-v3.5",
|
"rerank-v3.5",
|
||||||
|
|
@ -240,7 +196,6 @@ class TestGetReranker:
|
||||||
),
|
),
|
||||||
],
|
],
|
||||||
ids=[
|
ids=[
|
||||||
"mxbai",
|
|
||||||
"cohere",
|
"cohere",
|
||||||
"vllm",
|
"vllm",
|
||||||
"zeroentropy",
|
"zeroentropy",
|
||||||
|
|
|
||||||
38
uv.lock
38
uv.lock
|
|
@ -323,15 +323,6 @@ wheels = [
|
||||||
{ url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" },
|
{ url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "batched"
|
|
||||||
version = "0.1.5"
|
|
||||||
source = { registry = "https://pypi.org/simple" }
|
|
||||||
sdist = { url = "https://files.pythonhosted.org/packages/c9/40/8d9a8ed9b95cb95acf599698557b7074b462df652823a61e7e43899aa519/batched-0.1.5.tar.gz", hash = "sha256:58b8a41d3f8d4d39a0edba79c6238ed204938cfc2c8908224919d70af07c610d", size = 23940, upload-time = "2025-07-14T09:58:31.862Z" }
|
|
||||||
wheels = [
|
|
||||||
{ url = "https://files.pythonhosted.org/packages/fb/c8/16a977fd90cdc974ef7781e237b8a0e0008a6204768ededbef2b2ff1bb43/batched-0.1.5-py3-none-any.whl", hash = "sha256:356dae99f15c906629992e4bd3481a857114790b5316268fa38fe8ad0d0b9480", size = 29367, upload-time = "2025-07-14T09:58:30.968Z" },
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "beartype"
|
name = "beartype"
|
||||||
version = "0.22.9"
|
version = "0.22.9"
|
||||||
|
|
@ -1583,7 +1574,7 @@ name = "haiku-rag"
|
||||||
version = "0.65.1"
|
version = "0.65.1"
|
||||||
source = { editable = "." }
|
source = { editable = "." }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
|
{ name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "tui", "voyageai", "zeroentropy"] },
|
||||||
]
|
]
|
||||||
|
|
||||||
[package.optional-dependencies]
|
[package.optional-dependencies]
|
||||||
|
|
@ -1619,7 +1610,7 @@ dev = [
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
{ name = "haiku-rag-slim", extras = ["cross-encoder"], marker = "extra == 'cross-encoder'", editable = "haiku_rag_slim" },
|
{ name = "haiku-rag-slim", extras = ["cross-encoder"], marker = "extra == 'cross-encoder'", editable = "haiku_rag_slim" },
|
||||||
{ name = "haiku-rag-slim", extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "tui", "cross-encoder"], editable = "haiku_rag_slim" },
|
{ name = "haiku-rag-slim", extras = ["docling", "voyageai", "cohere", "zeroentropy", "tui", "cross-encoder"], editable = "haiku_rag_slim" },
|
||||||
{ name = "haiku-rag-slim", extras = ["ingester"], marker = "extra == 'ingester'", editable = "haiku_rag_slim" },
|
{ name = "haiku-rag-slim", extras = ["ingester"], marker = "extra == 'ingester'", editable = "haiku_rag_slim" },
|
||||||
{ name = "haiku-rag-slim", extras = ["s3"], marker = "extra == 's3'", editable = "haiku_rag_slim" },
|
{ name = "haiku-rag-slim", extras = ["s3"], marker = "extra == 's3'", editable = "haiku_rag_slim" },
|
||||||
{ name = "textual", marker = "extra == 'tui'", specifier = ">=8.2.4" },
|
{ name = "textual", marker = "extra == 'tui'", specifier = ">=8.2.4" },
|
||||||
|
|
@ -1732,10 +1723,6 @@ jina = [
|
||||||
mistral = [
|
mistral = [
|
||||||
{ name = "pydantic-ai-slim", extra = ["mistral"] },
|
{ name = "pydantic-ai-slim", extra = ["mistral"] },
|
||||||
]
|
]
|
||||||
mxbai = [
|
|
||||||
{ name = "mxbai-rerank" },
|
|
||||||
{ name = "transformers" },
|
|
||||||
]
|
|
||||||
s3 = [
|
s3 = [
|
||||||
{ name = "obstore" },
|
{ name = "obstore" },
|
||||||
]
|
]
|
||||||
|
|
@ -1769,7 +1756,6 @@ requires-dist = [
|
||||||
{ name = "jinja2", specifier = ">=3.1.0" },
|
{ name = "jinja2", specifier = ">=3.1.0" },
|
||||||
{ name = "jsonpatch", specifier = ">=1.33" },
|
{ name = "jsonpatch", specifier = ">=1.33" },
|
||||||
{ name = "lancedb", specifier = "==0.30.2" },
|
{ name = "lancedb", specifier = "==0.30.2" },
|
||||||
{ name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" },
|
|
||||||
{ name = "obstore", marker = "extra == 's3'", specifier = ">=0.9,<0.10" },
|
{ name = "obstore", marker = "extra == 's3'", specifier = ">=0.9,<0.10" },
|
||||||
{ name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.6.0.66,<5.0.0.0" },
|
{ name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.6.0.66,<5.0.0.0" },
|
||||||
{ name = "pathspec", specifier = ">=1.0.4" },
|
{ name = "pathspec", specifier = ">=1.0.4" },
|
||||||
|
|
@ -1793,7 +1779,6 @@ requires-dist = [
|
||||||
{ name = "textual-image", specifier = ">=0.8.5" },
|
{ name = "textual-image", specifier = ">=0.8.5" },
|
||||||
{ name = "torch", marker = "extra == 'jina'", specifier = ">=2.0.0" },
|
{ name = "torch", marker = "extra == 'jina'", specifier = ">=2.0.0" },
|
||||||
{ name = "transformers", marker = "extra == 'jina'", specifier = ">=4.40.0" },
|
{ name = "transformers", marker = "extra == 'jina'", specifier = ">=4.40.0" },
|
||||||
{ name = "transformers", marker = "extra == 'mxbai'", specifier = ">=4.49.0,<6.0.0" },
|
|
||||||
{ name = "tree-sitter", marker = "extra == 'tui'", specifier = ">=0.25.2" },
|
{ name = "tree-sitter", marker = "extra == 'tui'", specifier = ">=0.25.2" },
|
||||||
{ name = "tree-sitter-json", marker = "extra == 'tui'", specifier = ">=0.24.8" },
|
{ name = "tree-sitter-json", marker = "extra == 'tui'", specifier = ">=0.24.8" },
|
||||||
{ name = "typer", specifier = ">=0.21.0,<0.22.0" },
|
{ name = "typer", specifier = ">=0.21.0,<0.22.0" },
|
||||||
|
|
@ -1802,7 +1787,7 @@ requires-dist = [
|
||||||
{ name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a11" },
|
{ name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a11" },
|
||||||
{ name = "zstandard", marker = "python_full_version < '3.14'", specifier = ">=0.23.0" },
|
{ name = "zstandard", marker = "python_full_version < '3.14'", specifier = ">=0.23.0" },
|
||||||
]
|
]
|
||||||
provides-extras = ["docling", "s3", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "cross-encoder", "ingester", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
|
provides-extras = ["docling", "s3", "voyageai", "cohere", "zeroentropy", "jina", "cross-encoder", "ingester", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "haiku-skills"
|
name = "haiku-skills"
|
||||||
|
|
@ -2814,23 +2799,6 @@ wheels = [
|
||||||
{ url = "https://files.pythonhosted.org/packages/7e/82/69e539c4c2027f1e1697e09aaa2449243085a0edf81ae2c6341e84d769b6/multiprocess-0.70.19-py39-none-any.whl", hash = "sha256:0d4b4397ed669d371c81dcd1ef33fd384a44d6c3de1bd0ca7ac06d837720d3c5", size = 133477, upload-time = "2026-01-19T06:47:38.619Z" },
|
{ url = "https://files.pythonhosted.org/packages/7e/82/69e539c4c2027f1e1697e09aaa2449243085a0edf81ae2c6341e84d769b6/multiprocess-0.70.19-py39-none-any.whl", hash = "sha256:0d4b4397ed669d371c81dcd1ef33fd384a44d6c3de1bd0ca7ac06d837720d3c5", size = 133477, upload-time = "2026-01-19T06:47:38.619Z" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
|
||||||
name = "mxbai-rerank"
|
|
||||||
version = "0.1.6"
|
|
||||||
source = { registry = "https://pypi.org/simple" }
|
|
||||||
dependencies = [
|
|
||||||
{ name = "accelerate" },
|
|
||||||
{ name = "batched" },
|
|
||||||
{ name = "numpy" },
|
|
||||||
{ name = "torch" },
|
|
||||||
{ name = "tqdm" },
|
|
||||||
{ name = "transformers" },
|
|
||||||
]
|
|
||||||
sdist = { url = "https://files.pythonhosted.org/packages/0f/76/a19c864a1025222d3304a888ed4ed9217bfdf55dbaf4ed37500ee03935e0/mxbai_rerank-0.1.6.tar.gz", hash = "sha256:8d08e8464796429a7415314ce6de682bf9b538eb4ee5a7ddcd1a07839ee02879", size = 21449, upload-time = "2025-06-02T14:59:42.45Z" }
|
|
||||||
wheels = [
|
|
||||||
{ url = "https://files.pythonhosted.org/packages/a6/2a/503622b3a80272c662dabef421c9635168e5cbf6d51f0aa1883998561292/mxbai_rerank-0.1.6-py3-none-any.whl", hash = "sha256:aee94e7a14d5fba6520052ff2098f0f03db6cd9cc39553b7d2e82389deec9e05", size = 18458, upload-time = "2025-06-02T14:59:41.003Z" },
|
|
||||||
]
|
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "networkx"
|
name = "networkx"
|
||||||
version = "3.6.1"
|
version = "3.6.1"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue