feat(retrieval): pin model revisions and share encoder/reranker instances
Pin the Qwen embedding/reranker models to specific commit revisions (and sentence-transformers to the exact version already in use) so index builds are reproducible instead of floating on each model's main branch. Also move the encoder and reranker instances to module-level singletons shared by every Index, instead of one pair per instance, so a process serving multiple repository roots only loads each model once. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
1 parent
96a1175e20
commit
208213c1c7
3 files changed
+43
-14
No files matched your search
@@ -200,6 +200,18 @@ class RetrievalServerTests(unittest.TestCase):
|
|||||||
finally:
|
finally:
|
||||||
index.db.close()
|
index.db.close()
|
||||||
|
|
||||||
|
def test_shared_encoder_and_reranker_are_cached_module_singletons(self):
|
||||||
|
with patch.object(MODULE, '_shared_encoder', None), patch.object(MODULE, '_shared_reranker', None), \
|
||||||
|
patch.object(MODULE, 'preferred_device', lambda: 'cpu'):
|
||||||
|
sentinel = SimpleNamespace()
|
||||||
|
with patch.dict('sys.modules', sentence_transformers=SimpleNamespace(
|
||||||
|
SentenceTransformer=lambda *args, **kwargs: sentinel,
|
||||||
|
CrossEncoder=lambda *args, **kwargs: sentinel)):
|
||||||
|
self.assertIs(MODULE.shared_encoder(), sentinel)
|
||||||
|
self.assertIs(MODULE.shared_encoder(), MODULE.shared_encoder())
|
||||||
|
self.assertIs(MODULE.shared_reranker(), sentinel)
|
||||||
|
self.assertIs(MODULE.shared_reranker(), MODULE.shared_reranker())
|
||||||
|
|
||||||
def test_rerank_uses_cross_encoder_scores(self):
|
def test_rerank_uses_cross_encoder_scores(self):
|
||||||
index = MODULE.Index.__new__(MODULE.Index)
|
index = MODULE.Index.__new__(MODULE.Index)
|
||||||
|
|
||||||
|
|||||||
@@ -1 +1 @@
|
|||||||
sentence-transformers>=5.4,<6
|
sentence-transformers==5.7.0
|
||||||
+30
-13
@@ -18,7 +18,9 @@ import threading
|
|||||||
|
|
||||||
|
|
||||||
MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
MODEL = "Qwen/Qwen3-Embedding-0.6B"
|
||||||
|
MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3"
|
||||||
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
|
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
|
||||||
|
RERANKER_MODEL_REVISION = "e61197ed45024b0ed8a2d74b80b4d909f1255473"
|
||||||
EMBEDDING_DIMENSIONS = 1024
|
EMBEDDING_DIMENSIONS = 1024
|
||||||
CHUNK_TOKENS = 400
|
CHUNK_TOKENS = 400
|
||||||
OVERLAP_TOKENS = 64
|
OVERLAP_TOKENS = 64
|
||||||
@@ -163,6 +165,32 @@ def read_max_files():
|
|||||||
return parsed
|
return parsed
|
||||||
|
|
||||||
|
|
||||||
|
_shared_model_lock = threading.Lock()
|
||||||
|
_shared_encoder = None
|
||||||
|
_shared_reranker = None
|
||||||
|
|
||||||
|
|
||||||
|
def shared_encoder():
|
||||||
|
global _shared_encoder
|
||||||
|
with _shared_model_lock:
|
||||||
|
if _shared_encoder is None:
|
||||||
|
from sentence_transformers import SentenceTransformer
|
||||||
|
_shared_encoder = SentenceTransformer(MODEL, revision=MODEL_REVISION, truncate_dim=EMBEDDING_DIMENSIONS,
|
||||||
|
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
|
||||||
|
device=preferred_device())
|
||||||
|
return _shared_encoder
|
||||||
|
|
||||||
|
|
||||||
|
def shared_reranker():
|
||||||
|
global _shared_reranker
|
||||||
|
with _shared_model_lock:
|
||||||
|
if _shared_reranker is None:
|
||||||
|
from sentence_transformers import CrossEncoder
|
||||||
|
_shared_reranker = CrossEncoder(RERANKER_MODEL, revision=RERANKER_MODEL_REVISION, prompts={"query": RETRIEVAL_INSTRUCTION},
|
||||||
|
default_prompt_name="query", device=preferred_device())
|
||||||
|
return _shared_reranker
|
||||||
|
|
||||||
|
|
||||||
def reciprocal_rank_fusion(*rankings):
|
def reciprocal_rank_fusion(*rankings):
|
||||||
scores = {}
|
scores = {}
|
||||||
for ranking in rankings:
|
for ranking in rankings:
|
||||||
@@ -187,8 +215,6 @@ class Index:
|
|||||||
self.db.execute("create table if not exists query_cache (cache_key text primary key, result text not null)")
|
self.db.execute("create table if not exists query_cache (cache_key text primary key, result text not null)")
|
||||||
self.db.commit()
|
self.db.commit()
|
||||||
self.lock = threading.RLock()
|
self.lock = threading.RLock()
|
||||||
self.model = None
|
|
||||||
self.reranker_model = None
|
|
||||||
self.max_files = read_max_files()
|
self.max_files = read_max_files()
|
||||||
|
|
||||||
def files(self):
|
def files(self):
|
||||||
@@ -212,19 +238,10 @@ class Index:
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
def encoder(self):
|
def encoder(self):
|
||||||
if self.model is None:
|
return shared_encoder()
|
||||||
from sentence_transformers import SentenceTransformer
|
|
||||||
self.model = SentenceTransformer(MODEL, truncate_dim=EMBEDDING_DIMENSIONS,
|
|
||||||
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
|
|
||||||
device=preferred_device())
|
|
||||||
return self.model
|
|
||||||
|
|
||||||
def reranker(self):
|
def reranker(self):
|
||||||
if self.reranker_model is None:
|
return shared_reranker()
|
||||||
from sentence_transformers import CrossEncoder
|
|
||||||
self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION},
|
|
||||||
default_prompt_name="query", device=preferred_device())
|
|
||||||
return self.reranker_model
|
|
||||||
|
|
||||||
def rerank(self, query: str, rows: list[dict], limit: int):
|
def rerank(self, query: str, rows: list[dict], limit: int):
|
||||||
candidates = rows[:RERANK_CANDIDATES]
|
candidates = rows[:RERANK_CANDIDATES]
|
||||||
|
|||||||
Reference in new issue
Block a user