refac

backend/open_webui/retrieval/models/colbert.py

@@ -13,19 +13,17 @@ log = logging.getLogger(__name__)
 
 class ColBERT(BaseReranker):
     def __init__(self, name, **kwargs) -> None:
-        log.info("ColBERT: Loading model", name)
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        log.info('ColBERT: Loading model', name)
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-        DOCKER = kwargs.get("env") == "docker"
+        DOCKER = kwargs.get('env') == 'docker'
         if DOCKER:
             # This is a workaround for the issue with the docker container
             # where the torch extension is not loaded properly
             # and the following error is thrown:
             # /root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/segmented_maxsim_cpp.so: cannot open shared object file: No such file or directory
 
-            lock_file = (
-                "/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock"
-            )
+            lock_file = '/root/.cache/torch_extensions/py311_cpu/segmented_maxsim_cpp/lock'
             if os.path.exists(lock_file):
                 os.remove(lock_file)
 
@@ -36,23 +34,16 @@ class ColBERT(BaseReranker):
         pass
 
     def calculate_similarity_scores(self, query_embeddings, document_embeddings):
-
         query_embeddings = query_embeddings.to(self.device)
         document_embeddings = document_embeddings.to(self.device)
 
         # Validate dimensions to ensure compatibility
         if query_embeddings.dim() != 3:
-            raise ValueError(
-                f"Expected query embeddings to have 3 dimensions, but got {query_embeddings.dim()}."
-            )
+            raise ValueError(f'Expected query embeddings to have 3 dimensions, but got {query_embeddings.dim()}.')
         if document_embeddings.dim() != 3:
-            raise ValueError(
-                f"Expected document embeddings to have 3 dimensions, but got {document_embeddings.dim()}."
-            )
+            raise ValueError(f'Expected document embeddings to have 3 dimensions, but got {document_embeddings.dim()}.')
         if query_embeddings.size(0) not in [1, document_embeddings.size(0)]:
-            raise ValueError(
-                "There should be either one query or queries equal to the number of documents."
-            )
+            raise ValueError('There should be either one query or queries equal to the number of documents.')
 
         # Transpose the query embeddings to align for matrix multiplication
         transposed_query_embeddings = query_embeddings.permute(0, 2, 1)
@@ -69,7 +60,6 @@ class ColBERT(BaseReranker):
         return normalized_scores.detach().cpu().numpy().astype(np.float32)
 
     def predict(self, sentences):
-
         query = sentences[0][0]
         docs = [i[1] for i in sentences]
 
@@ -80,8 +70,6 @@ class ColBERT(BaseReranker):
         embedded_query = embedded_queries[0]
 
         # Calculate retrieval scores for the query against all documents
-        scores = self.calculate_similarity_scores(
-            embedded_query.unsqueeze(0), embedded_docs
-        )
+        scores = self.calculate_similarity_scores(embedded_query.unsqueeze(0), embedded_docs)
 
         return scores