comments

Author: jihan-yin

nucleus/__init__.py

@@ -1105,6 +1105,15 @@ def delete_model(self, model_id: str) -> dict:
     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_urls,

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]