updated OPQ to reduce dimensions to 4 times of output subvectors

beir-cellar · Aug 3, 2023 · 1ea9387 · 1ea9387
1 parent 9fb620f
commit 1ea9387
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/beir/retrieval/search/dense/faiss_search.py b/beir/retrieval/search/dense/faiss_search.py
@@ -196,7 +196,9 @@ def index(self, corpus: Dict[str, Dict[str, str]], score_function: str = None, *
         if self.use_rotation:
             logger.info("Rotating data before encoding it with a product quantizer...")
             logger.info("Creating OPQ Matrix...")
-            opq_matrix = faiss.OPQMatrix(self.dim_size, self.code_size)
+            logger.info("Input Dimension: {}, Output Dimension: {}".format(self.dim_size, self.num_of_centroids*4))
+            opq_matrix = faiss.OPQMatrix(self.dim_size, self.code_size, self.num_of_centroids*4)
+            base_index = faiss.IndexPQ(self.num_of_centroids*4, self.num_of_centroids, self.code_size, self.similarity_metric)
             base_index = faiss.IndexPreTransform(opq_matrix, base_index)
 
         if self.use_gpu: