Remove Ollama reranker, this was a bad idea
This commit is contained in:
parent
e08fce84f8
commit
fe850e778d
6 changed files with 12 additions and 126 deletions
|
|
@ -15,11 +15,11 @@ The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 fo
|
||||||
| Embedding Model | Document in top 1 | Document in top 3 | Reranker |
|
| Embedding Model | Document in top 1 | Document in top 3 | Reranker |
|
||||||
|---------------------------------------|-------------------|-------------------|------------------------|
|
|---------------------------------------|-------------------|-------------------|------------------------|
|
||||||
| Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None |
|
| Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None |
|
||||||
| Ollama / `mxbai-embed-large` | 0.81 | 0.91 | `mxbai-rerank-base-v2` |
|
| Ollama / `mxbai-embed-large` | 0.90 | 0.95 | `mxbai-rerank-base-v2` |
|
||||||
| Ollama / `nomic-embed-text` | 0.74 | 0.88 | None |
|
<!-- | Ollama / `nomic-embed-text` | 0.74 | 0.88 | None |
|
||||||
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
|
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
|
||||||
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
|
| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | None |
|
||||||
| OpenAI / `text-embeddings-3-small` | 0.83 | 0.90 | Cohere / `rerank-v3.5` |
|
| OpenAI / `text-embeddings-3-small` | 0.83 | 0.90 | Cohere / `rerank-v3.5` | -->
|
||||||
|
|
||||||
## Question/Answer evaluation
|
## Question/Answer evaluation
|
||||||
|
|
||||||
|
|
@ -29,5 +29,5 @@ Again using the same dataset, we use a QA agent to answer the question. In addit
|
||||||
|------------------------------------|-----------------------------------|-----------|------------------------|
|
|------------------------------------|-----------------------------------|-----------|------------------------|
|
||||||
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.85 | None |
|
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.85 | None |
|
||||||
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.72 | `mxbai-rerank-base-v2` |
|
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.72 | `mxbai-rerank-base-v2` |
|
||||||
| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None |
|
<!-- | Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None |
|
||||||
| OpenAI / `text-embeddings-3-small` | OpenAI / `gpt-4-turbo` | 0.62 | None |
|
| OpenAI / `text-embeddings-3-small` | OpenAI / `gpt-4-turbo` | 0.62 | None | -->
|
||||||
|
|
|
||||||
|
|
@ -109,25 +109,7 @@ See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the com
|
||||||
|
|
||||||
Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results.
|
Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results.
|
||||||
|
|
||||||
Reranking is **disabled by default** for faster searches. You can enable it by configuring a reranking provider.
|
Reranking is **disabled by default** (`RERANK_PROVIDER=""`) for faster searches. You can enable it by configuring one of the providers below.
|
||||||
|
|
||||||
### Disabling Reranking
|
|
||||||
|
|
||||||
To disable reranking completely for faster searches:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
RERANK_PROVIDER=""
|
|
||||||
```
|
|
||||||
|
|
||||||
### Ollama (Default)
|
|
||||||
|
|
||||||
Ollama reranking uses LLMs with structured output to rank documents by relevance:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
RERANK_PROVIDER="ollama"
|
|
||||||
RERANK_MODEL="qwen3:1.7b" # or any model that supports structured output
|
|
||||||
OLLAMA_BASE_URL="http://localhost:11434"
|
|
||||||
```
|
|
||||||
|
|
||||||
### MixedBread AI
|
### MixedBread AI
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -20,7 +20,7 @@ class AppConfig(BaseModel):
|
||||||
EMBEDDINGS_VECTOR_DIM: int = 1024
|
EMBEDDINGS_VECTOR_DIM: int = 1024
|
||||||
|
|
||||||
RERANK_PROVIDER: str = ""
|
RERANK_PROVIDER: str = ""
|
||||||
RERANK_MODEL: str = "qwen3"
|
RERANK_MODEL: str = ""
|
||||||
|
|
||||||
QA_PROVIDER: str = "ollama"
|
QA_PROVIDER: str = "ollama"
|
||||||
QA_MODEL: str = "qwen3"
|
QA_MODEL: str = "qwen3"
|
||||||
|
|
|
||||||
|
|
@ -31,10 +31,4 @@ def get_reranker() -> RerankerBase | None:
|
||||||
except ImportError:
|
except ImportError:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
if Config.RERANK_PROVIDER == "ollama":
|
|
||||||
from haiku.rag.reranking.ollama import OllamaReranker
|
|
||||||
|
|
||||||
_reranker = OllamaReranker()
|
|
||||||
return _reranker
|
|
||||||
|
|
||||||
return None
|
return None
|
||||||
|
|
|
||||||
|
|
@ -1,81 +0,0 @@
|
||||||
from pydantic import BaseModel
|
|
||||||
from pydantic_ai import Agent
|
|
||||||
from pydantic_ai.models.openai import OpenAIModel
|
|
||||||
from pydantic_ai.providers.ollama import OllamaProvider
|
|
||||||
|
|
||||||
from haiku.rag.config import Config
|
|
||||||
from haiku.rag.reranking.base import RerankerBase
|
|
||||||
from haiku.rag.store.models.chunk import Chunk
|
|
||||||
|
|
||||||
|
|
||||||
class RerankResult(BaseModel):
|
|
||||||
"""Individual rerank result with index and relevance score."""
|
|
||||||
|
|
||||||
index: int
|
|
||||||
relevance_score: float
|
|
||||||
|
|
||||||
|
|
||||||
class RerankResponse(BaseModel):
|
|
||||||
"""Response from the reranking model containing ranked results."""
|
|
||||||
|
|
||||||
results: list[RerankResult]
|
|
||||||
|
|
||||||
|
|
||||||
class OllamaReranker(RerankerBase):
|
|
||||||
def __init__(self, model: str = Config.RERANK_MODEL):
|
|
||||||
self._model = model
|
|
||||||
|
|
||||||
# Create the reranking prompt
|
|
||||||
system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query.
|
|
||||||
|
|
||||||
Return your response as a JSON object with a "results" array. Each result should have:
|
|
||||||
- "index": the original index of the document (integer)
|
|
||||||
- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant)
|
|
||||||
|
|
||||||
Only return the top documents up to the requested limit, ordered by decreasing relevance score.
|
|
||||||
/no_think
|
|
||||||
"""
|
|
||||||
|
|
||||||
model_obj = OpenAIModel(
|
|
||||||
model_name=model,
|
|
||||||
provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
|
|
||||||
)
|
|
||||||
|
|
||||||
self._agent = Agent(
|
|
||||||
model=model_obj,
|
|
||||||
output_type=RerankResponse,
|
|
||||||
system_prompt=system_prompt,
|
|
||||||
)
|
|
||||||
|
|
||||||
async def rerank(
|
|
||||||
self, query: str, chunks: list[Chunk], top_n: int = 10
|
|
||||||
) -> list[tuple[Chunk, float]]:
|
|
||||||
if not chunks:
|
|
||||||
return []
|
|
||||||
|
|
||||||
documents = []
|
|
||||||
for i, chunk in enumerate(chunks):
|
|
||||||
documents.append({"index": i, "content": chunk.content})
|
|
||||||
|
|
||||||
documents_text = ""
|
|
||||||
for doc in documents:
|
|
||||||
documents_text += f"Index {doc['index']}: {doc['content']}\n\n"
|
|
||||||
|
|
||||||
user_prompt = f"""Query: {query}
|
|
||||||
|
|
||||||
Documents to rerank:
|
|
||||||
{documents_text.strip()}
|
|
||||||
|
|
||||||
Rank these documents by relevance to the query and return the top {top_n} results as JSON."""
|
|
||||||
|
|
||||||
try:
|
|
||||||
result = await self._agent.run(user_prompt)
|
|
||||||
|
|
||||||
return [
|
|
||||||
(chunks[result_item.index], result_item.relevance_score)
|
|
||||||
for result_item in result.output.results[:top_n]
|
|
||||||
]
|
|
||||||
|
|
||||||
except Exception:
|
|
||||||
# Fallback: return chunks in original order with same score
|
|
||||||
return [(chunks[i], 1.0) for i in range(min(top_n, len(chunks)))]
|
|
||||||
|
|
@ -24,7 +24,7 @@ chunks = [
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reranker_base():
|
async def test_reranker_base():
|
||||||
reranker = RerankerBase()
|
reranker = RerankerBase()
|
||||||
assert reranker._model == "qwen3"
|
assert reranker._model == ""
|
||||||
|
|
||||||
with pytest.raises(NotImplementedError):
|
with pytest.raises(NotImplementedError):
|
||||||
await reranker.rerank("query", [])
|
await reranker.rerank("query", [])
|
||||||
|
|
@ -35,12 +35,16 @@ async def test_mxbai_reranker():
|
||||||
try:
|
try:
|
||||||
from haiku.rag.reranking.mxbai import MxBAIReranker
|
from haiku.rag.reranking.mxbai import MxBAIReranker
|
||||||
|
|
||||||
|
Config.RERANK_MODEL = "mixedbread-ai/mxbai-rerank-base-v2"
|
||||||
reranker = MxBAIReranker()
|
reranker = MxBAIReranker()
|
||||||
|
# reranker._model = "mixedbread-ai/mxbai-rerank-base-v2"
|
||||||
reranked = await reranker.rerank(
|
reranked = await reranker.rerank(
|
||||||
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
||||||
)
|
)
|
||||||
assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
|
assert [chunk.document_id for chunk, score in reranked] == ["0", "2"]
|
||||||
assert all(isinstance(score, float) for chunk, score in reranked)
|
assert all(isinstance(score, float) for chunk, score in reranked)
|
||||||
|
Config.RERANK_MODEL = ""
|
||||||
|
|
||||||
except ImportError:
|
except ImportError:
|
||||||
pytest.skip("MxBAI package not installed")
|
pytest.skip("MxBAI package not installed")
|
||||||
|
|
||||||
|
|
@ -62,16 +66,3 @@ async def test_cohere_reranker():
|
||||||
|
|
||||||
except ImportError:
|
except ImportError:
|
||||||
pytest.skip("Cohere package not installed")
|
pytest.skip("Cohere package not installed")
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
|
||||||
async def test_ollama_reranker():
|
|
||||||
from haiku.rag.reranking.ollama import OllamaReranker
|
|
||||||
|
|
||||||
reranker = OllamaReranker()
|
|
||||||
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)
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue