[WIP] Add muvera 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,32 @@ def __add_props(
         ret_dict["properties"] = existing_props
 
 
+class _VectorIndexMultivectorEncoding:
+    @staticmethod
+    def muvera(
+        enabled: bool = True,
+        ksim: Optional[int] = None,
+        dprojections: Optional[int] = None,
+        repetitions: Optional[int] = None,
+    ) -> _EncodingConfigCreate:
+        return _MuveraConfigCreate(
+            enabled=enabled,
+            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("encoding") is None
+        or not config.get("encoding", {"enabled": False}).get("enabled")
+        else _MuveraConfig(
+            enabled=config["encoding"]["enabled"],
+            ksim=config["encoding"]["ksim"],
+            dprojections=config["encoding"]["dprojections"],
+            repetitions=config["encoding"]["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("encoding") is None
+                else __get_multivector_encoding(config["multivector"]["encoding"])
+            ),
             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