diff --git a/remote_vector_index_builder/core/common/models/index_builder/faiss/ivf_pq_build_cagra_config.py b/remote_vector_index_builder/core/common/models/index_builder/faiss/ivf_pq_build_cagra_config.py index 1d34674..c645217 100755 --- a/remote_vector_index_builder/core/common/models/index_builder/faiss/ivf_pq_build_cagra_config.py +++ b/remote_vector_index_builder/core/common/models/index_builder/faiss/ivf_pq_build_cagra_config.py @@ -53,6 +53,8 @@ class IVFPQBuildCagraConfig: # to use as little GPU memory for the database as possible. conservative_memory_allocation: bool = True + force_random_rotation: bool = False + @staticmethod def _validate_params(params: Dict[str, Any]) -> None: """ @@ -120,6 +122,7 @@ def to_faiss_config(self) -> faiss.IVFPQBuildCagraConfig: config.pq_dim = self.pq_dim config.n_lists = self.n_lists config.conservative_memory_allocation = self.conservative_memory_allocation + config.force_random_rotation = self.force_random_rotation return config @classmethod diff --git a/test_remote_vector_index_builder/test_core/common/models/index_builder/faiss/test_ivf_pq_build_cagra_config.py b/test_remote_vector_index_builder/test_core/common/models/index_builder/faiss/test_ivf_pq_build_cagra_config.py index 7936fe1..8977299 100644 --- a/test_remote_vector_index_builder/test_core/common/models/index_builder/faiss/test_ivf_pq_build_cagra_config.py +++ b/test_remote_vector_index_builder/test_core/common/models/index_builder/faiss/test_ivf_pq_build_cagra_config.py @@ -27,6 +27,7 @@ def custom_params(self) -> Dict[str, Any]: "pq_bits": 6, "pq_dim": 16, "conservative_memory_allocation": False, + "force_random_rotation": True, } def test_default_initialization(self, default_config): @@ -36,6 +37,7 @@ def test_default_initialization(self, default_config): assert default_config.pq_bits == 8 assert default_config.pq_dim == 0 assert default_config.conservative_memory_allocation is True + assert default_config.force_random_rotation is False def test_custom_initialization(self, custom_params): config = IVFPQBuildCagraConfig(**custom_params) @@ -97,6 +99,7 @@ def test_to_faiss_config(self, custom_params): config.conservative_memory_allocation == custom_params["conservative_memory_allocation"] ) + assert config.force_random_rotation == custom_params["force_random_rotation"] def test_from_dict_partial(self): partial_params = {"n_lists": 2048, "kmeans_n_iters": 30} @@ -107,3 +110,4 @@ def test_from_dict_partial(self): assert config.pq_bits == 8 # default value assert config.pq_dim == 0 # default value assert config.conservative_memory_allocation is True # default value + assert config.force_random_rotation is False