Allow dimensions to be passed in config, update tests

weaviate · Nov 27, 2024 · 75ffad5 · 75ffad5
1 parent 9c6178c
commit 75ffad5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/test/collection/test_config.py b/test/collection/test_config.py
@@ -402,12 +402,14 @@ def test_basic_config():
             vectorize_collection_name=False,
             model="Snowflake/snowflake-arctic-embed-m-v1.5",
             base_url="https://api.embedding.weaviate.io",
+            dimensions=768,
         ),
         {
             "text2vec-weaviate": {
                 "vectorizeClassName": False,
                 "model": "Snowflake/snowflake-arctic-embed-m-v1.5",
                 "baseURL": "https://api.embedding.weaviate.io",
+                "dimensions": 768,
             }
         },
     ),
@@ -1508,6 +1510,7 @@ def test_vector_config_flat_pq() -> None:
                 name="test",
                 source_properties=["prop"],
                 base_url="https://api.embedding.weaviate.io",
+                dimensions=768,
             )
         ],
         {
@@ -1517,6 +1520,7 @@ def test_vector_config_flat_pq() -> None:
                         "properties": ["prop"],
                         "vectorizeClassName": True,
                         "baseURL": "https://api.embedding.weaviate.io",
+                        "dimensions": 768,
                     }
                 },
                 "vectorIndexType": "hnsw",

diff --git a/weaviate/collections/classes/config_named_vectors.py b/weaviate/collections/classes/config_named_vectors.py
@@ -1228,6 +1228,7 @@ def text2vec_weaviate(
         vectorize_collection_name: bool = True,
         model: Optional[Union[WeaviateModel, str]] = None,
         base_url: Optional[str] = None,
+        dimensions: Optional[int] = None,
     ) -> _NamedVectorConfigCreate:
         return _NamedVectorConfigCreate(
             name=name,
@@ -1236,6 +1237,7 @@ def text2vec_weaviate(
                 model=model,
                 vectorizeClassName=vectorize_collection_name,
                 baseURL=base_url,
+                dimensions=dimensions,
             ),
             vector_index_config=vector_index_config,
         )

diff --git a/weaviate/collections/classes/config_vectorizers.py b/weaviate/collections/classes/config_vectorizers.py
@@ -351,6 +351,7 @@ class _Text2VecWeaviateConfig(_VectorizerConfigCreate):
     model: Optional[str]
     baseURL: Optional[str]
     vectorizeClassName: bool
+    dimensions: Optional[int]
 
 
 class _Text2VecOllamaConfig(_VectorizerConfigCreate):
@@ -1302,10 +1303,12 @@ def text2vec_weaviate(
         model: Optional[Union[WeaviateModel, str]] = None,
         base_url: Optional[str] = None,
         vectorize_collection_name: bool = True,
+        dimensions: Optional[int] = None,
     ) -> _VectorizerConfigCreate:
         """TODO: add docstrings when the documentation is available."""
         return _Text2VecWeaviateConfig(
             model=model,
             baseURL=base_url,
             vectorizeClassName=vectorize_collection_name,
+            dimensions=dimensions,
         )