Merge pull request #90 from weaviate/rodrigo/vector_dimensions

ingest_data: Add vector dimensions as a parameter

test/unittests/test_managers/test_data_manager.py

@@ -24,6 +24,7 @@ def test_ingest_data(mock_client):
         consistency_level="quorum",
         randomize=True,
         auto_tenants=0,
+        vector_dimensions=1536,
     )
 
     mock_client.collections.get.assert_called_once_with("TestCollection")

weaviate_cli/commands/create.py

@@ -237,10 +237,29 @@ def create_backup_cli(ctx, backend, backup_id, include, exclude, wait, cpu_for_b
     default=0,
     help="Number of tenants for which we will send data. NOTE: Requires class with --auto_tenant_creation (default: 0).",
 )
+@click.option(
+    "--vector_dimensions",
+    default=1536,
+    help="Number of vector dimensions to be used when the data is randomized.",
+)
 @click.pass_context
-def create_data_cli(ctx, collection, limit, consistency_level, randomize, auto_tenants):
+def create_data_cli(
+    ctx,
+    collection,
+    limit,
+    consistency_level,
+    randomize,
+    auto_tenants,
+    vector_dimensions,
+):
     """Ingest data into a collection in Weaviate."""
 
+    if vector_dimensions != 1536 and not randomize:
+        click.echo(
+            "Error: --vector_dimensions has no effect unless --randomize is enabled."
+        )
+        sys.exit(1)
+
     client = None
     try:
         client = get_client_from_context(ctx)
@@ -252,6 +271,7 @@ def create_data_cli(ctx, collection, limit, consistency_level, randomize, auto_t
             consistency_level=consistency_level,
             randomize=randomize,
             auto_tenants=auto_tenants,
+            vector_dimensions=vector_dimensions,
         )
     except Exception as e:
         click.echo(f"Error: {e}")

weaviate_cli/managers/data_manager.py

@@ -124,22 +124,34 @@ def __ingest_data(
         num_objects: int,
         cl: wvc.ConsistencyLevel,
         randomize: bool,
+        vector_dimensions: Optional[int] = 1536,
     ) -> int:
         if randomize:
             counter = 0
             data_objects = self.__generate_data_object(num_objects)
             cl_collection = collection.with_consistency_level(cl)
             vectorizer = cl_collection.config.get().vectorizer
-            dimensions = 1536
             if vectorizer == "text2vec-contextionary":
-                dimensions = 300
+                (
+                    print("Warning: Using vector dimensions: 300")
+                    if vector_dimensions != 1536
+                    else None
+                )
+                vector_dimensions = 300
             elif vectorizer == "text2vec-transformers":
-                dimensions = 768
+                (
+                    print("Warning: Using vector dimensions: 768")
+                    if vector_dimensions != 1536
+                    else None
+                )
+                vector_dimensions = 768
             with cl_collection.batch.dynamic() as batch:
                 for obj in data_objects:
                     batch.add_object(
                         properties=obj,
-                        vector=np.random.rand(1, dimensions)[0].tolist(),
+                        vector=(
+                            2 * np.random.rand(1, vector_dimensions)[0] - 1
+                        ).tolist(),
                     )
                     counter += 1
 
@@ -167,6 +179,7 @@ def ingest_data(
         consistency_level: str,
         randomize: bool,
         auto_tenants: int,
+        vector_dimensions: Optional[int] = 1536,
     ) -> None:
 
         if not self.client.collections.exists(collection):
@@ -218,6 +231,7 @@ def ingest_data(
                     limit,
                     cl_map[consistency_level],
                     randomize,
+                    vector_dimensions,
                 )
             else:
                 click.echo(f"Processing tenant '{tenant}'")
@@ -226,6 +240,7 @@ def ingest_data(
                     limit,
                     cl_map[consistency_level],
                     randomize,
+                    vector_dimensions,
                 )
 
             if ret == -1: