Support for jina reranker, both local and API
This commit is contained in:
parent
b8cf8f5198
commit
7ff014684a
8 changed files with 308 additions and 1 deletions
|
|
@ -382,3 +382,43 @@ reranking:
|
||||||
```
|
```
|
||||||
|
|
||||||
**Note:** vLLM reranking uses the `/v1/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded.
|
**Note:** vLLM reranking uses the `/v1/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded.
|
||||||
|
|
||||||
|
### Jina AI
|
||||||
|
|
||||||
|
Jina provides high-quality reranking with two deployment options: API mode and local inference.
|
||||||
|
|
||||||
|
#### API Mode
|
||||||
|
|
||||||
|
Use the Jina Reranker API for cloud-based reranking:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: jina
|
||||||
|
name: jina-reranker-v3
|
||||||
|
```
|
||||||
|
|
||||||
|
Set your API key via environment variable:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
export JINA_API_KEY=your-api-key
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Local Mode
|
||||||
|
|
||||||
|
For local inference, install the jina extra:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv pip install haiku.rag-slim[jina]
|
||||||
|
```
|
||||||
|
|
||||||
|
Then configure:
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
reranking:
|
||||||
|
model:
|
||||||
|
provider: jina-local
|
||||||
|
name: jinaai/jina-reranker-v3
|
||||||
|
```
|
||||||
|
|
||||||
|
**Note:** The Jina Reranker v3 local model is licensed under CC BY-NC 4.0, which restricts commercial use. For commercial applications, use the API mode instead.
|
||||||
|
|
|
||||||
|
|
@ -52,4 +52,19 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
except ImportError: # pragma: no cover
|
except ImportError: # pragma: no cover
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
if config.reranking.model and config.reranking.model.provider == "jina":
|
||||||
|
from haiku.rag.reranking.jina import JinaReranker
|
||||||
|
|
||||||
|
model = config.reranking.model.name or "jina-reranker-v3"
|
||||||
|
return JinaReranker(model)
|
||||||
|
|
||||||
|
if config.reranking.model and config.reranking.model.provider == "jina-local":
|
||||||
|
try:
|
||||||
|
from haiku.rag.reranking.jina_local import JinaLocalReranker
|
||||||
|
|
||||||
|
model = config.reranking.model.name or "jinaai/jina-reranker-v3"
|
||||||
|
return JinaLocalReranker(model)
|
||||||
|
except ImportError: # pragma: no cover
|
||||||
|
return None
|
||||||
|
|
||||||
return None
|
return None
|
||||||
|
|
|
||||||
50
haiku_rag_slim/haiku/rag/reranking/jina.py
Normal file
50
haiku_rag_slim/haiku/rag/reranking/jina.py
Normal file
|
|
@ -0,0 +1,50 @@
|
||||||
|
import os
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
|
||||||
|
from haiku.rag.reranking.base import RerankerBase
|
||||||
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
|
class JinaReranker(RerankerBase):
|
||||||
|
"""Jina AI reranker using the Jina Reranker API."""
|
||||||
|
|
||||||
|
def __init__(self, model: str = "jina-reranker-v3"):
|
||||||
|
self._model = model
|
||||||
|
self._api_key = os.environ.get("JINA_API_KEY")
|
||||||
|
if not self._api_key:
|
||||||
|
raise ValueError("JINA_API_KEY environment variable required")
|
||||||
|
|
||||||
|
async def rerank(
|
||||||
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
|
) -> list[tuple[Chunk, float]]:
|
||||||
|
if not chunks:
|
||||||
|
return []
|
||||||
|
|
||||||
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
|
||||||
|
async with httpx.AsyncClient() as client:
|
||||||
|
response = await client.post(
|
||||||
|
"https://api.jina.ai/v1/rerank",
|
||||||
|
json={
|
||||||
|
"model": self._model,
|
||||||
|
"query": query,
|
||||||
|
"documents": documents,
|
||||||
|
"top_n": top_n,
|
||||||
|
},
|
||||||
|
headers={
|
||||||
|
"Authorization": f"Bearer {self._api_key}",
|
||||||
|
"Content-Type": "application/json",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
|
||||||
|
result = response.json()
|
||||||
|
|
||||||
|
scored_chunks = []
|
||||||
|
for item in result.get("results", []):
|
||||||
|
index = item["index"]
|
||||||
|
score = item["relevance_score"]
|
||||||
|
scored_chunks.append((chunks[index], score))
|
||||||
|
|
||||||
|
return scored_chunks
|
||||||
46
haiku_rag_slim/haiku/rag/reranking/jina_local.py
Normal file
46
haiku_rag_slim/haiku/rag/reranking/jina_local.py
Normal file
|
|
@ -0,0 +1,46 @@
|
||||||
|
try:
|
||||||
|
from transformers import (
|
||||||
|
AutoModelForSequenceClassification, # pyright: ignore[reportMissingImports]
|
||||||
|
)
|
||||||
|
except ImportError as e:
|
||||||
|
raise ImportError(
|
||||||
|
"transformers is not installed. Please install it with `pip install transformers torch` "
|
||||||
|
"or use the jina optional dependency."
|
||||||
|
) from e
|
||||||
|
|
||||||
|
from haiku.rag.reranking.base import RerankerBase
|
||||||
|
from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
|
|
||||||
|
class JinaLocalReranker(RerankerBase):
|
||||||
|
"""Jina reranker using local model inference via transformers.
|
||||||
|
|
||||||
|
Note: The Jina Reranker v3 model is licensed under CC BY-NC 4.0,
|
||||||
|
which restricts commercial use.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, model: str = "jinaai/jina-reranker-v3"):
|
||||||
|
self._model = model
|
||||||
|
self._reranker = AutoModelForSequenceClassification.from_pretrained(
|
||||||
|
model, trust_remote_code=True
|
||||||
|
)
|
||||||
|
|
||||||
|
async def rerank(
|
||||||
|
self, query: str, chunks: list[Chunk], top_n: int = 10
|
||||||
|
) -> list[tuple[Chunk, float]]:
|
||||||
|
if not chunks:
|
||||||
|
return []
|
||||||
|
|
||||||
|
documents = [chunk.content for chunk in chunks]
|
||||||
|
sentence_pairs = [[query, doc] for doc in documents]
|
||||||
|
|
||||||
|
scores = self._reranker.compute_score(sentence_pairs)
|
||||||
|
|
||||||
|
# Handle both single score and list of scores
|
||||||
|
if isinstance(scores, (int, float)):
|
||||||
|
scores = [scores]
|
||||||
|
|
||||||
|
scored_chunks = list(zip(chunks, scores, strict=False))
|
||||||
|
scored_chunks.sort(key=lambda x: x[1], reverse=True)
|
||||||
|
|
||||||
|
return [(chunk, float(score)) for chunk, score in scored_chunks[:top_n]]
|
||||||
|
|
@ -44,6 +44,7 @@ voyageai = ["voyageai>=0.3.7"]
|
||||||
mxbai = ["mxbai-rerank>=0.1.6"]
|
mxbai = ["mxbai-rerank>=0.1.6"]
|
||||||
cohere = ["cohere>=5.20.1"]
|
cohere = ["cohere>=5.20.1"]
|
||||||
zeroentropy = ["zeroentropy>=0.1.0a7"]
|
zeroentropy = ["zeroentropy>=0.1.0a7"]
|
||||||
|
jina = ["transformers>=4.40.0", "torch>=2.0.0"]
|
||||||
# TUI (chat and inspect commands)
|
# TUI (chat and inspect commands)
|
||||||
tui = ["textual>=7.3.0", "textual-image>=0.8.5"]
|
tui = ["textual>=7.3.0", "textual-image>=0.8.5"]
|
||||||
# Model providers (delegated to pydantic-ai-slim)
|
# Model providers (delegated to pydantic-ai-slim)
|
||||||
|
|
|
||||||
76
tests/cassettes/test_reranker/test_jina_reranker.yaml
Normal file
76
tests/cassettes/test_reranker/test_jina_reranker.yaml
Normal file
|
|
@ -0,0 +1,76 @@
|
||||||
|
interactions:
|
||||||
|
- request:
|
||||||
|
headers:
|
||||||
|
accept:
|
||||||
|
- '*/*'
|
||||||
|
accept-encoding:
|
||||||
|
- gzip, deflate, zstd
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
content-length:
|
||||||
|
- '1266'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
host:
|
||||||
|
- api.jina.ai
|
||||||
|
method: POST
|
||||||
|
parsed_body:
|
||||||
|
documents:
|
||||||
|
- To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer
|
||||||
|
Prize, and has become a classic of modern American literature.
|
||||||
|
- The novel Moby-Dick was written by Herman Melville and first published in 1851. It is considered a masterpiece of
|
||||||
|
American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.
|
||||||
|
- Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville,
|
||||||
|
Alabama. She received the Pulitzer Prize for Fiction in 1961.
|
||||||
|
- Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment
|
||||||
|
upon the British landed gentry at the end of the 18th century.
|
||||||
|
- The Harry Potter series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the
|
||||||
|
most popular and critically acclaimed books of the modern era.
|
||||||
|
- The Great Gatsby, a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set
|
||||||
|
in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan.
|
||||||
|
model: jina-reranker-v3
|
||||||
|
query: Who wrote 'To Kill a Mockingbird'?
|
||||||
|
top_n: 2
|
||||||
|
uri: https://api.jina.ai/v1/rerank
|
||||||
|
response:
|
||||||
|
headers:
|
||||||
|
alt-svc:
|
||||||
|
- h3=":443"; ma=86400
|
||||||
|
cache-control:
|
||||||
|
- private
|
||||||
|
connection:
|
||||||
|
- keep-alive
|
||||||
|
content-length:
|
||||||
|
- '570'
|
||||||
|
content-type:
|
||||||
|
- application/json
|
||||||
|
expires:
|
||||||
|
- Wed, 21 Jan 2026 08:21:12 GMT
|
||||||
|
nel:
|
||||||
|
- '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
|
||||||
|
report-to:
|
||||||
|
- '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=TX0p0eBpMlJ8P0KA1iF7CLglvZSaumFNwZXrU%2BYgFDRyrC0xrRB3NyaekMeKtZQKWf8GpGD%2Bb5WCh4vN67uaCzgah3QbrVh3PvU%2FSEMF6BZcfFMY"}]}'
|
||||||
|
transfer-encoding:
|
||||||
|
- chunked
|
||||||
|
vary:
|
||||||
|
- Accept-Encoding
|
||||||
|
parsed_body:
|
||||||
|
model: jina-reranker-v3
|
||||||
|
object: list
|
||||||
|
results:
|
||||||
|
- document:
|
||||||
|
text: To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the
|
||||||
|
Pulitzer Prize, and has become a classic of modern American literature.
|
||||||
|
index: 0
|
||||||
|
relevance_score: 0.42858967
|
||||||
|
- document:
|
||||||
|
text: Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville,
|
||||||
|
Alabama. She received the Pulitzer Prize for Fiction in 1961.
|
||||||
|
index: 2
|
||||||
|
relevance_score: 0.07394931
|
||||||
|
usage:
|
||||||
|
total_tokens: 490
|
||||||
|
status:
|
||||||
|
code: 200
|
||||||
|
message: OK
|
||||||
|
version: 1
|
||||||
|
|
@ -202,3 +202,76 @@ class TestGetReranker:
|
||||||
)
|
)
|
||||||
result = get_reranker(config)
|
result = get_reranker(config)
|
||||||
assert result is None
|
assert result is None
|
||||||
|
|
||||||
|
def test_jina_provider(self, monkeypatch):
|
||||||
|
monkeypatch.setenv("JINA_API_KEY", "test-api-key")
|
||||||
|
|
||||||
|
from haiku.rag.reranking.jina import JinaReranker
|
||||||
|
|
||||||
|
config = AppConfig(
|
||||||
|
reranking=RerankingConfig(
|
||||||
|
model=ModelConfig(provider="jina", name="jina-reranker-v3")
|
||||||
|
)
|
||||||
|
)
|
||||||
|
result = get_reranker(config)
|
||||||
|
assert isinstance(result, JinaReranker)
|
||||||
|
assert result._model == "jina-reranker-v3"
|
||||||
|
|
||||||
|
def test_jina_local_provider(self):
|
||||||
|
try:
|
||||||
|
from haiku.rag.reranking.jina_local import JinaLocalReranker
|
||||||
|
|
||||||
|
config = AppConfig(
|
||||||
|
reranking=RerankingConfig(
|
||||||
|
model=ModelConfig(
|
||||||
|
provider="jina-local", name="jinaai/jina-reranker-v3"
|
||||||
|
)
|
||||||
|
)
|
||||||
|
)
|
||||||
|
result = get_reranker(config)
|
||||||
|
assert isinstance(result, JinaLocalReranker)
|
||||||
|
assert result._model == "jinaai/jina-reranker-v3"
|
||||||
|
except ImportError:
|
||||||
|
pytest.skip("Jina local dependencies not installed")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
@pytest.mark.vcr()
|
||||||
|
async def test_jina_reranker(monkeypatch):
|
||||||
|
import os
|
||||||
|
|
||||||
|
# Only set dummy key if real key not present (for VCR playback)
|
||||||
|
if not os.environ.get("JINA_API_KEY"):
|
||||||
|
monkeypatch.setenv("JINA_API_KEY", "test-api-key")
|
||||||
|
|
||||||
|
from haiku.rag.reranking.jina import JinaReranker
|
||||||
|
|
||||||
|
reranker = JinaReranker("jina-reranker-v3")
|
||||||
|
|
||||||
|
reranked = await reranker.rerank(
|
||||||
|
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
||||||
|
)
|
||||||
|
assert len(reranked) == 2
|
||||||
|
assert all(isinstance(score, float) for chunk, score in reranked)
|
||||||
|
# Check that the top results are relevant to Harper Lee / To Kill a Mockingbird
|
||||||
|
top_ids = [chunk.document_id for chunk, score in reranked]
|
||||||
|
assert "0" in top_ids or "2" in top_ids # These chunks mention the book/author
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_jina_local_reranker():
|
||||||
|
try:
|
||||||
|
from haiku.rag.reranking.jina_local import JinaLocalReranker
|
||||||
|
|
||||||
|
reranker = JinaLocalReranker("jinaai/jina-reranker-v3")
|
||||||
|
|
||||||
|
reranked = await reranker.rerank(
|
||||||
|
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
|
||||||
|
)
|
||||||
|
assert len(reranked) == 2
|
||||||
|
assert all(isinstance(score, float) for chunk, score in reranked)
|
||||||
|
# Check that the top results are relevant to Harper Lee / To Kill a Mockingbird
|
||||||
|
top_ids = [chunk.document_id for chunk, score in reranked]
|
||||||
|
assert "0" in top_ids or "2" in top_ids # These chunks mention the book/author
|
||||||
|
except ImportError:
|
||||||
|
pytest.skip("Jina local dependencies not installed")
|
||||||
|
|
|
||||||
8
uv.lock
8
uv.lock
|
|
@ -1397,6 +1397,10 @@ google = [
|
||||||
groq = [
|
groq = [
|
||||||
{ name = "pydantic-ai-slim", extra = ["groq"] },
|
{ name = "pydantic-ai-slim", extra = ["groq"] },
|
||||||
]
|
]
|
||||||
|
jina = [
|
||||||
|
{ name = "torch" },
|
||||||
|
{ name = "transformers" },
|
||||||
|
]
|
||||||
mistral = [
|
mistral = [
|
||||||
{ name = "pydantic-ai-slim", extra = ["mistral"] },
|
{ name = "pydantic-ai-slim", extra = ["mistral"] },
|
||||||
]
|
]
|
||||||
|
|
@ -1440,12 +1444,14 @@ requires-dist = [
|
||||||
{ name = "rich", specifier = ">=14.2.0" },
|
{ name = "rich", specifier = ">=14.2.0" },
|
||||||
{ name = "textual", marker = "extra == 'tui'", specifier = ">=7.3.0" },
|
{ name = "textual", marker = "extra == 'tui'", specifier = ">=7.3.0" },
|
||||||
{ name = "textual-image", marker = "extra == 'tui'", specifier = ">=0.8.5" },
|
{ 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 = "typer", specifier = ">=0.19.2,<0.20.0" },
|
{ name = "typer", specifier = ">=0.19.2,<0.20.0" },
|
||||||
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.7" },
|
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.7" },
|
||||||
{ name = "watchfiles", specifier = ">=1.1.1" },
|
{ name = "watchfiles", specifier = ">=1.1.1" },
|
||||||
{ name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a7" },
|
{ name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a7" },
|
||||||
]
|
]
|
||||||
provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
|
provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "hf-xet"
|
name = "hf-xet"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue