From 208213c1c7e9cbefa980afdd48a057d252c73804 Mon Sep 17 00:00:00 2001 From: Julian lechner Date: Fri, 18 Sep 2026 09:35:46 +0200 Subject: [PATCH] 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 --- scripts/test-retrieval-server.py | 12 +++++++++ shared/retrieval/requirements.txt | 2 +- shared/retrieval/server.py | 43 +++++++++++++++++++++---------- 3 files changed, 43 insertions(+), 14 deletions(-) diff --git a/scripts/test-retrieval-server.py b/scripts/test-retrieval-server.py index bb6bbc8..cea0722 100644 --- a/scripts/test-retrieval-server.py +++ b/scripts/test-retrieval-server.py @@ -200,6 +200,18 @@ class RetrievalServerTests(unittest.TestCase): finally: 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): index = MODULE.Index.__new__(MODULE.Index) diff --git a/shared/retrieval/requirements.txt b/shared/retrieval/requirements.txt index e842278..28cee65 100644 --- a/shared/retrieval/requirements.txt +++ b/shared/retrieval/requirements.txt @@ -1 +1 @@ -sentence-transformers>=5.4,<6 +sentence-transformers==5.7.0 diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 303857e..932ac48 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -18,7 +18,9 @@ import threading MODEL = "Qwen/Qwen3-Embedding-0.6B" +MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3" RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B" +RERANKER_MODEL_REVISION = "e61197ed45024b0ed8a2d74b80b4d909f1255473" EMBEDDING_DIMENSIONS = 1024 CHUNK_TOKENS = 400 OVERLAP_TOKENS = 64 @@ -163,6 +165,32 @@ def read_max_files(): 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): scores = {} 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.commit() self.lock = threading.RLock() - self.model = None - self.reranker_model = None self.max_files = read_max_files() def files(self): @@ -212,19 +238,10 @@ class Index: return result def encoder(self): - if self.model is None: - 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 + return shared_encoder() def reranker(self): - if self.reranker_model is None: - 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 + return shared_reranker() def rerank(self, query: str, rows: list[dict], limit: int): candidates = rows[:RERANK_CANDIDATES]