Merge pull request #85 from scaleapi/jihan/custom-efs

Nucleus - custom indexing improvements

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,23 @@ 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
+    ):
+        """
+        Creates a custom index for a given dataset, which will then be used
+        for autotag and similarity search.
+
+        :param
+        dataset_id: id of dataset that the custom index is being added to.
+        embeddings_urls: list of urls, each of which being a json mapping dataset_item_id -> embedding vector
+        embedding_dim: the dimension of the embedding vectors, must be consistent for all embedding vectors in the index.
+        """
         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)

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.1.13"
+version = "0.1.14"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/helpers.py

@@ -14,11 +14,11 @@
 
 
 TEST_IMG_URLS = [
-    "https://homepages.cae.wisc.edu/~ece533/images/airplane.png",
-    "https://homepages.cae.wisc.edu/~ece533/images/arctichare.png",
-    "https://homepages.cae.wisc.edu/~ece533/images/baboon.png",
-    "https://homepages.cae.wisc.edu/~ece533/images/barbara.png",
-    "https://homepages.cae.wisc.edu/~ece533/images/cat.png",
+    "https://d3jkudlc7u70kh.cloudfront.net/airplane-fact.jpg",
+    "https://www.activewild.com/wp-content/uploads/2016/01/Arctic-Hare-Facts-For-Kids-fb.jpg",
+    "https://thumbs-prod.si-cdn.com/Gd4SBf-ePT2OK6HDdW6Q3uOh5c8=/fit-in/1600x0/https://public-media.si-cdn.com/filer/23/a9/23a9a987-df8e-4d22-ac96-93f9091f9b8f/media_baboons_07_-_photo_by_catherine_markham.jpg",
+    "https://prod-images.tcm.com/Master-Profile-Images/BarbaraStanwyck.jpg",
+    "https://thumbs-prod.si-cdn.com/nnJARGtKrLypH4y3Vov2zGTG4xw=/fit-in/1600x0/filters:focal(554x699:555x700)/https://public-media.si-cdn.com/filer/a4/04/a404c799-7118-459a-8de4-89e4a44b124f/img_1317.jpg",
 ]
 
 TEST_DATASET_ITEMS = [

tests/test_indexing.py

@@ -39,7 +39,9 @@ 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], embedding_dim=3
+    )
     assert JOB_ID_KEY in create_response
     assert MESSAGE_KEY in create_response
     job_id = create_response[JOB_ID_KEY]