feat(retrieval): prefer CUDA devices

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']

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):