add gpu mapping

src/retrieval/context.py

@@ -12,6 +12,9 @@ import torch
 import numpy as np
 from qdrant_client.http import models as rest
 
+# Import device detection utility
+from src.utils.device import get_device_for_sentence_transformers
+
 
 try:
     from langchain.docstore.document import Document
@@ -57,10 +60,12 @@ class ContextRetriever:
                 from colbert.infra import Run, ColBERTConfig
                 from colbert.modeling.checkpoint import Checkpoint
                 # ColBERT uses late interaction - different implementation needed
+                device = get_device_for_sentence_transformers()
                 print(f"✅ RERANKER: ColBERT model detected ({self.reranker_model_name})")
                 print(f"🔍 INTERACTION TYPE: Late interaction (token-level embeddings)")
+                print(f"🖥️  DEVICE: {device}")
 
-                # Create ColBERT config for CPU mode
+                # Create ColBERT config with device
                 colbert_config = ColBERTConfig(
                     doc_maxlen=300,
                     query_maxlen=32,
@@ -72,15 +77,18 @@ class ContextRetriever:
                 # Load checkpoint (e.g. "colbert-ir/colbertv2.0")
                 self.colbert_checkpoint = Checkpoint(self.reranker_model_name, colbert_config=colbert_config)
                 self.colbert_model = self.colbert_checkpoint.model
+                # Move model to device
+                self.colbert_model = self.colbert_model.to(device)
                 self.colbert_tokenizer = self.colbert_checkpoint.raw_tokenizer
                 self.reranker = self._colbert_rerank  # attach wrapper function
-                print(f"✅ COLBERT: Model and tokenizer loaded successfully")
+                print(f"✅ COLBERT: Model and tokenizer loaded successfully on {device}")
 
             else:
                 # Standard CrossEncoder for BGE and other models
                 from sentence_transformers import CrossEncoder
-                self.reranker = CrossEncoder(self.reranker_model_name)
-                print(f"✅ RERANKER: Initialized {self.reranker_model_name}")
+                device = get_device_for_sentence_transformers()
+                self.reranker = CrossEncoder(self.reranker_model_name, device=device)
+                print(f"✅ RERANKER: Initialized {self.reranker_model_name} on {device}")
                 print(f"🔍 INTERACTION TYPE: Cross-encoder (single relevance score)")
         except Exception as e:
             print(f"⚠️ Reranker initialization failed: {e}")