support custom indexing with efs

Author: jihan-yin

nucleus/__init__.py

@@ -82,6 +82,7 @@
     DATASET_ITEM_IDS_KEY,
     DEFAULT_NETWORK_TIMEOUT_SEC,
     EMBEDDINGS_URL_KEY,
+    EMBEDDING_DIMENSION_KEY,
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
@@ -1101,9 +1102,14 @@ def delete_model(self, model_id: str) -> dict:
         )
         return response
 
-    def create_custom_index(self, dataset_id: str, embeddings_url: str):
+    def create_custom_index(
+        self, dataset_id: str, embeddings_urls: list, embedding_dim: int
+    ):
         return self.make_request(
-            {EMBEDDINGS_URL_KEY: embeddings_url},
+            {
+                EMBEDDINGS_URL_KEY: embeddings_urls,
+                EMBEDDING_DIMENSION_KEY: embedding_dim,
+            },
             f"indexing/{dataset_id}",
             requests_command=requests.post,
         )

nucleus/constants.py

@@ -24,7 +24,8 @@
 DEFAULT_ANNOTATION_UPDATE_MODE = False
 DEFAULT_NETWORK_TIMEOUT_SEC = 120
 DIMENSIONS_KEY = "dimensions"
-EMBEDDINGS_URL_KEY = "embeddings_url"
+EMBEDDINGS_URL_KEY = "embeddings_urls"
+EMBEDDING_DIMENSION_KEY = "embedding_dimension"
 ERRORS_KEY = "errors"
 ERROR_CODES = "error_codes"
 ERROR_ITEMS = "upload_errors"

nucleus/dataset.py

@@ -329,8 +329,12 @@ def delete_item(self, item_id: str = None, reference_id: str = None):
     def list_autotags(self):
         return self._client.list_autotags(self.id)
 
-    def create_custom_index(self, embeddings_url: str):
-        return self._client.create_custom_index(self.id, embeddings_url)
+    def create_custom_index(self, embeddings_urls: list, embedding_dim: int):
+        return self._client.create_custom_index(
+            self.id,
+            embeddings_urls,
+            embedding_dim,
+        )
 
     def delete_custom_index(self):
         return self._client.delete_custom_index(self.id)

tests/test_indexing.py

@@ -39,7 +39,7 @@ def dataset(CLIENT):
 
 def test_index_integration(dataset):
     signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
-    create_response = dataset.create_custom_index(signed_embeddings_url)
+    create_response = dataset.create_custom_index([signed_embeddings_url])
     assert JOB_ID_KEY in create_response
     assert MESSAGE_KEY in create_response
     job_id = create_response[JOB_ID_KEY]