Fix vector DB benchmarks to use model_inputs configuration (#470)

* Initial plan

* Fix vector DB benchmarks to use model_inputs configuration

Co-authored-by: drewoldag <47493171+drewoldag@users.noreply.github.com>

* Fix Qdrant vector_size configuration to match actual vector length

Co-authored-by: drewoldag <47493171+drewoldag@users.noreply.github.com>

* Slight modification to the model_inputs definition.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: drewoldag <47493171+drewoldag@users.noreply.github.com>

Author: Copilot

benchmarks/vector_db_benchmarks.py

@@ -24,7 +24,13 @@ def setup(self, vector_length, vector_db_implementation):
 
         self.h = Hyrax()
         self.h.config["general"]["results_dir"] = str(self.input_dir)
-        self.h.config["data_set"]["name"] = "HyraxRandomDataset"
+        self.h.config["model_inputs"] = {
+            "data": {
+                "dataset_class": "HyraxRandomDataset",
+                "fields": ["image", "label", "object_id"],
+                "primary_id_field": "object_id",
+            }
+        }
         self.h.config["model"]["name"] = "HyraxLoopback"
 
         # Default inference batch size is 512, so this should result in 4 batch files
@@ -83,7 +89,13 @@ def setup(self, shard_size_limit, vector_db_implementation):
 
         self.h = Hyrax()
         self.h.config["general"]["results_dir"] = str(self.input_dir)
-        self.h.config["data_set"]["name"] = "HyraxRandomDataset"
+        self.h.config["model_inputs"] = {
+            "data": {
+                "dataset_class": "HyraxRandomDataset",
+                "fields": ["image", "label", "object_id"],
+                "primary_id_field": "object_id",
+            }
+        }
         self.h.config["data_loader"]["batch_size"] = 4096
         self.h.config["model"]["name"] = "HyraxLoopback"
 
@@ -102,12 +114,12 @@ def setup(self, shard_size_limit, vector_db_implementation):
 
         # Get the list of dataset ids
         self.ds = self.h.prepare()
-        self.data_sample = self.ds[4001]["data"]["image"].numpy()
+        self.data_sample = self.ds[0]["data"]["image"]
 
         self.h.config["vector_db"]["name"] = vector_db_implementation
         self.h.config["vector_db"]["chromadb"]["shard_size_limit"] = shard_size_limit
         # Qdrant requires the vector size in order to create its collections
-        self.h.config["vector_db"]["qdrant"]["vector_size"] = 4096
+        self.h.config["vector_db"]["qdrant"]["vector_size"] = self.vector_length
 
         # Save inference results to vector database and create a db connection
         self.h.save_to_database(output_dir=Path(self.output_dir))