Add a new subset to show it on HF data viewer

Signed-off-by: Jay Wang <jay@zijie.wang>
poloclub · Jan 22, 2024 · 52c2349 · 52c2349
1 parent 864881c
commit 52c2349
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/scripts/diffusiondb.py b/scripts/diffusiondb.py
@@ -232,6 +232,18 @@ class DiffusionDB(datasets.GeneratorBasedBuilder):
         ),
     )
 
+    # Add a random 1k from 2M as the first entry point to show on HF data viewer
+    # Sample part_ids
+    part_ids = np.random.choice(_PART_IDS, 1000, replace=False).tolist()
+    BUILDER_CONFIGS.append(
+        DiffusionDBConfig(
+            name="1k_random_2m",
+            part_ids=part_ids,
+            is_large=False,
+            description="Another random 1k images with meta data from DiffusionDB 2M",
+        ),
+    )
+
     # Default to only load 1k random images
     DEFAULT_CONFIG_NAME = "2m_random_1k"