danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Jina uses AutoModel not AutoModelForSequenceClassification

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-21 11:26:23 +02:00
parent ce95dc47a5
commit a1abfd9666
No known key found for this signature in database
2 changed files with 7 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
haiku_rag_slim/haiku/rag/client.py
View file

@ -1715,11 +1715,11 @@ class HaikuRAG:
elif provider == "jina-local":
try:
from transformers import AutoModelForSequenceClassification
from transformers import AutoModel
yield DownloadProgress(model=model_name, status="start")
await asyncio.to_thread(
AutoModelForSequenceClassification.from_pretrained,
AutoModel.from_pretrained,
model_name,
trust_remote_code=True,
)

19
haiku_rag_slim/haiku/rag/reranking/jina_local.py
View file

@ -1,6 +1,6 @@
try:
from transformers import (
AutoModelForSequenceClassification, # pyright: ignore[reportMissingImports]
AutoModel, # pyright: ignore[reportMissingImports]
)
except ImportError as e:
raise ImportError(
@ -21,9 +21,8 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover
def __init__(self, model: str = "jinaai/jina-reranker-v3"):
self._model = model
self._reranker = AutoModelForSequenceClassification.from_pretrained(
model, trust_remote_code=True
)
self._reranker = AutoModel.from_pretrained(model, trust_remote_code=True)
self._reranker.eval()
async def rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
@ -32,15 +31,7 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover
return []
documents = [chunk.content for chunk in chunks]
sentence_pairs = [[query, doc] for doc in documents]
scores = self._reranker.compute_score(sentence_pairs)
results = self._reranker.rerank(query, documents, top_n=top_n)
# 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]]
return [(chunks[r["index"]], float(r["relevance_score"])) for r in results]