Fix BFloat16 numpy conversion

visual_rag/retrieval/multi_vector.py

@@ -157,9 +157,10 @@ class MultiVectorRetriever:
     ) -> List[Dict[str, Any]]:
         q = self.embedder.embed_query(query)
         if isinstance(q, torch.Tensor):
-            query_embedding = q.detach().cpu().numpy()
+            # .float() converts BFloat16 to Float32 (numpy doesn't support BFloat16)
+            query_embedding = q.detach().cpu().float().numpy()
         else:
-            query_embedding = np.asarray(q)
+            query_embedding = np.asarray(q, dtype=np.float32)
 
         return self.search_embedded(
             query_embedding=query_embedding,

visual_rag/retrieval/single_stage.py

@@ -129,5 +129,5 @@ class SingleStageRetriever:
         if isinstance(embedding, torch.Tensor):
             if embedding.dtype == torch.bfloat16:
                 return embedding.cpu().float().numpy()
-            return embedding.cpu().numpy()
+            return embedding.cpu().float().numpy()  # .float() for BFloat16 compatibility
         return np.array(embedding, dtype=np.float32)

visual_rag/retrieval/three_stage.py

@@ -51,7 +51,7 @@ class ThreeStageRetriever:
         if isinstance(embedding, torch.Tensor):
             if embedding.dtype == torch.bfloat16:
                 return embedding.cpu().float().numpy()
-            return embedding.cpu().numpy()
+            return embedding.cpu().float().numpy()  # .float() for BFloat16 compatibility
         return np.array(embedding, dtype=np.float32)
 
     def _infer_vector_is_multivector(self, vector_name: str) -> bool:

visual_rag/retrieval/two_stage.py

@@ -418,7 +418,7 @@ class TwoStageRetriever:
         if isinstance(embedding, torch.Tensor):
             if embedding.dtype == torch.bfloat16:
                 return embedding.cpu().float().numpy()
-            return embedding.cpu().numpy()
+            return embedding.cpu().float().numpy()  # .float() for BFloat16 compatibility
         return np.array(embedding, dtype=np.float32)
 
     def build_filter(