From e5b63965b8ef497458b154a25783cce009dd5483 Mon Sep 17 00:00:00 2001 From: Julian lechner Date: Thu, 17 Sep 2026 17:50:33 +0200 Subject: [PATCH] feat(retrieval): prefer CUDA devices --- scripts/test-retrieval-server.py | 6 ++++++ shared/retrieval/server.py | 11 +++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/scripts/test-retrieval-server.py b/scripts/test-retrieval-server.py index 814e284..184decb 100644 --- a/scripts/test-retrieval-server.py +++ b/scripts/test-retrieval-server.py @@ -27,6 +27,12 @@ def embedding(): class RetrievalServerTests(unittest.TestCase): + def test_preferred_device_uses_cuda_with_cpu_fallback(self): + with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))): + self.assertEqual(MODULE.preferred_device(), 'cuda') + with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))): + self.assertEqual(MODULE.preferred_device(), 'cpu') + @unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.') def test_real_models(self): os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE'] diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index fef66e5..0d66e25 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -118,6 +118,11 @@ def document_text(path: str, symbol: str, content: str): return f"Path: {path}\nSymbol: {symbol}\n\n{content}" +def preferred_device(): + import torch + return "cuda" if torch.cuda.is_available() else "cpu" + + def reciprocal_rank_fusion(*rankings): scores = {} for ranking in rankings: @@ -165,13 +170,15 @@ class Index: 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: "}) + prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "}, + device=preferred_device()) return self.model 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") + 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):