limit datum upload to 30k

build_map.py

@@ -115,7 +115,8 @@ def upload_dataset_to_atlas(dataset_dict,
                             indexed_field = None, 
                             modality=None,
                             organization_name=None,
-                            wait_for_map=True):
+                            wait_for_map=True,
+                            datum_limit=30000):
 
     if modality is None:
         modality = "text"
@@ -124,7 +125,7 @@ def upload_dataset_to_atlas(dataset_dict,
         unique_id_field_name = "atlas_datum_id"
 
     if project_name is None:
-        project_name = dataset_dict["name"].replace("/", "--")
+        project_name = dataset_dict["name"].replace("/", "--") + "--hf-atlas-map"
 
     desc = f"Config: {dataset_dict['config']}"
 
@@ -169,13 +170,22 @@ def upload_dataset_to_atlas(dataset_dict,
     batch_size = 1000
     batched_texts = []
 
+    allow_upload = True
+
     for split in dataset_dict["splits"]:
 
+        if not allow_upload:
+            break
+
         dataset = load_dataset(dataset_dict["name"], dataset_dict["config"], split = split, streaming=True)
 
         for i, ex in tqdm(enumerate(dataset)):
             if i % 10000 == 0:
                 time.sleep(2)
+            if i == datum_limit:
+                print("Datum upload limited to 30,000 points. Stopping upload...")
+                allow_upload = False
+                break
 
             data_to_add = {"split": split, unique_id_field_name: f"{split}_{i}"}
 
@@ -245,4 +255,4 @@ if __name__ == "__main__":
     project_name = "huggingface_auto_upload_test-dolly-15k"
 
     dataset_dict = load_dataset_and_metadata(dataset_name)
-    print(upload_dataset_to_atlas(dataset_dict, project_name=project_name))
+    print(upload_dataset_to_atlas(dataset_dict, project_name=project_name))