Merge pull request #1635 from weaviate/rob/muvera

Add muvera config

Author: robbespo00

integration/test_collection_config.py

@@ -1405,6 +1405,67 @@ def test_config_multi_vector_disabled(
     assert conf.multi_vector is None
 
 
+def test_config_muvera_enabled(
+    collection_factory: CollectionFactory,
+) -> None:
+    dummy = collection_factory("dummy", ports=(8086, 50057))
+    if dummy._connection._weaviate_version.is_lower_than(1, 31, 0):
+        pytest.skip("Muvera is not supported in Weaviate versions lower than 1.31.0")
+
+    collection = collection_factory(
+        ports=(8086, 50057),
+        properties=[Property(name="name", data_type=DataType.TEXT)],
+        vectorizer_config=[
+            Configure.NamedVectors.text2colbert_jinaai(
+                name="vec",
+                vectorize_collection_name=False,
+                vector_index_config=Configure.VectorIndex.hnsw(
+                    multi_vector=Configure.VectorIndex.MultiVector.multi_vector(
+                        encoding=Configure.VectorIndex.MultiVector.Encoding.muvera()
+                    )
+                ),
+            )
+        ],
+    )
+    config = collection.config.get()
+    assert config.vector_config is not None
+    conf = config.vector_config["vec"].vector_index_config
+    assert isinstance(conf, _VectorIndexConfigHNSW)
+    if collection._connection._weaviate_version.is_lower_than(1, 31, 0):
+        assert conf.multi_vector is None
+    else:
+        assert conf.multi_vector is not None
+        assert conf.multi_vector.encoding is not None
+
+
+def test_config_muvera_disabled(
+    collection_factory: CollectionFactory,
+) -> None:
+    dummy = collection_factory("dummy", ports=(8086, 50057))
+    if dummy._connection._weaviate_version.is_lower_than(1, 29, 0):
+        pytest.skip("Multivector is not supported in Weaviate versions lower than 1.29.0")
+
+    collection = collection_factory(
+        ports=(8086, 50057),
+        properties=[Property(name="name", data_type=DataType.TEXT)],
+        vectorizer_config=[
+            Configure.NamedVectors.text2colbert_jinaai(
+                name="vec",
+                vectorize_collection_name=False,
+                vector_index_config=Configure.VectorIndex.hnsw(
+                    multi_vector=Configure.VectorIndex.MultiVector.multi_vector()
+                ),
+            )
+        ],
+    )
+    config = collection.config.get()
+    assert config.vector_config is not None
+    conf = config.vector_config["vec"].vector_index_config
+    assert isinstance(conf, _VectorIndexConfigHNSW)
+    assert conf.multi_vector is not None
+    assert conf.multi_vector.encoding is None
+
+
 @pytest.mark.parametrize(
     "generative_config",
     [

weaviate/collections/classes/config.py

@@ -33,6 +33,8 @@
 )
 from weaviate.collections.classes.config_vector_index import (
     VectorFilterStrategy,
+    _MuveraConfigCreate,
+    _EncodingConfigCreate,
     _MultiVectorConfigCreate,
     _QuantizerConfigCreate,
     _VectorIndexConfigCreate,
@@ -1556,8 +1558,20 @@ class _SQConfig(_ConfigBase):
 SQConfig = _SQConfig
 
 
+@dataclass
+class _MuveraConfig(_ConfigBase):
+    enabled: Optional[bool]
+    ksim: Optional[int]
+    dprojections: Optional[int]
+    repetitions: Optional[int]
+
+
+MuveraConfig = _MuveraConfig
+
+
 @dataclass
 class _MultiVectorConfig(_ConfigBase):
+    encoding: Optional[_MuveraConfig]
     aggregation: str
 
 
@@ -2031,12 +2045,31 @@ def __add_props(
         ret_dict["properties"] = existing_props
 
 
+class _VectorIndexMultivectorEncoding:
+    @staticmethod
+    def muvera(
+        ksim: Optional[int] = None,
+        dprojections: Optional[int] = None,
+        repetitions: Optional[int] = None,
+    ) -> _EncodingConfigCreate:
+        return _MuveraConfigCreate(
+            enabled=True,
+            ksim=ksim,
+            dprojections=dprojections,
+            repetitions=repetitions,
+        )
+
+
 class _VectorIndexMultiVector:
+    Encoding = _VectorIndexMultivectorEncoding
+
     @staticmethod
     def multi_vector(
+        encoding: Optional[_EncodingConfigCreate] = None,
         aggregation: Optional[MultiVectorAggregation] = None,
     ) -> _MultiVectorConfigCreate:
         return _MultiVectorConfigCreate(
+            encoding=encoding if encoding is not None else None,
             aggregation=aggregation.value if aggregation is not None else None,
         )
 

weaviate/collections/classes/config_methods.py

@@ -21,6 +21,7 @@
     _InvertedIndexConfig,
     _MultiTenancyConfig,
     _MultiVectorConfig,
+    _MuveraConfig,
     _NamedVectorConfig,
     _NamedVectorizerConfig,
     _NestedProperty,
@@ -149,12 +150,31 @@ def __get_quantizer_config(
     return quantizer
 
 
+def __get_multivector_encoding(config: Dict[str, Any]) -> Optional[_MuveraConfig]:
+    return (
+        None
+        if config.get("muvera") is None
+        or not config.get("muvera", {"enabled": False}).get("enabled")
+        else _MuveraConfig(
+            enabled=config["muvera"]["enabled"],
+            ksim=config["muvera"]["ksim"],
+            dprojections=config["muvera"]["dprojections"],
+            repetitions=config["muvera"]["repetitions"],
+        )
+    )
+
+
 def __get_multivector(config: Dict[str, Any]) -> Optional[_MultiVectorConfig]:
     return (
         None
         if config.get("multivector") is None
         or not config.get("multivector", {"enabled": False}).get("enabled")
         else _MultiVectorConfig(
+            encoding=(
+                None
+                if config["multivector"].get("muvera") is None
+                else __get_multivector_encoding(config["multivector"])
+            ),
             aggregation=config["multivector"]["aggregation"],
         )
     )
@@ -244,6 +264,7 @@ def __get_vector_config(
                 ),
                 vector_index_config=vector_index_config,
             )
+
         return named_vectors
     else:
         return None

weaviate/collections/classes/config_vector_index.py

@@ -42,7 +42,26 @@ class _MultiVectorConfigCreateBase(_ConfigCreateModel):
     enabled: bool = Field(default=True)
 
 
+class _EncodingConfigCreate(_MultiVectorConfigCreateBase):
+    enabled: bool = Field(default=True)
+
+    @staticmethod
+    @abstractmethod
+    def encoding_name() -> str: ...
+
+
+class _MuveraConfigCreate(_EncodingConfigCreate):
+    ksim: Optional[int]
+    dprojections: Optional[int]
+    repetitions: Optional[int]
+
+    @staticmethod
+    def encoding_name() -> str:
+        return "muvera"
+
+
 class _MultiVectorConfigCreate(_MultiVectorConfigCreateBase):
+    encoding: Optional[_EncodingConfigCreate] = Field(exclude=True)
     aggregation: Optional[str]
 
 
@@ -61,6 +80,10 @@ def _to_dict(self) -> Dict[str, Any]:
             ret_dict[self.quantizer.quantizer_name()] = self.quantizer._to_dict()
         if self.distance is not None:
             ret_dict["distance"] = str(self.distance.value)
+        if self.multivector is not None and self.multivector.encoding is not None:
+            ret_dict["multivector"][self.multivector.encoding.encoding_name()] = (
+                self.multivector.encoding._to_dict()
+            )
 
         return ret_dict