Adding new code to randomly sample from the observed distribution over labels.

app.py

@@ -1,8 +1,14 @@
 import os
+import json
+
 import gradio as gr
 import random
 import time
+
+from collections import Counter
+from numpy.random import choice
 from datasets import load_dataset, Dataset
+
 from PIL import PngImagePlugin, ImageFile
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 PngImagePlugin.MAX_TEXT_CHUNK = 1048576 * 10 # this is 10x the amount.
@@ -16,9 +22,11 @@ https://huggingface.co/datasets/datacomp
 # The number of examples/instances in this dataset is copied from the model card:
 # https://huggingface.co/datasets/ILSVRC/imagenet-1k
 NUM_EXAMPLES = 1281167
-# Arbitrary small number, only using in devv'ing (uncomment #.take(DEV_AMOUNT) below to use it).
-DEV = False
+# Arbitrary small number of dataset examples to look at, only using in devv'ing.
+DEV = True
 DEV_AMOUNT = 10
+# Whether to read in the distribution over labels from an external text file.
+READ_DISTRO = False
 GATED_IMAGENET = os.environ.get("GATED_IMAGENET")
 
 
@@ -41,10 +49,17 @@ def main(percentage=10):
         dataset = load_dataset("ILSVRC/imagenet-1k", split="train", streaming=True,
                            trust_remote_code=True, token=GATED_IMAGENET)
 
+    label_list, label_fractions = get_label_fractions()
+
+    # How many new random labels are we creating?
+    num_new_labels = int(round(NUM_EXAMPLES/float(percentage)))            
+    
     # Create a set of indices that are randomly chosen, to change their labels.
-    # Specifically, randomly choose NUM_EXAMPLES/percentage indices.
-    randomize_subset = set(random.sample(range(0, NUM_EXAMPLES), round(
-        NUM_EXAMPLES / float(percentage))))
+    # Specifically, randomly choose num_new_labels indices.
+    randomize_subset = set(random.sample(range(0, NUM_EXAMPLES), num_new_labels))
+    
+    # Randomly choose what the new label values are, following the observed label frequencies.
+    new_random_labels = list(choice(a=label_keys, size=num_new_labels, p=label_fractions))
 
     # Update the dataset so that the labels are randomized
     updated_dataset = dataset.map(randomize_labels, with_indices=True,
@@ -53,7 +68,7 @@ def main(percentage=10):
     # Upload the new version of the dataset (this will take awhile)
     if DEV:
         Dataset.from_generator(updated_dataset.__iter__).push_to_hub(
-            "datacomp/imagenet-1k-random-debug" + str(percentage), token=GATED_IMAGENET)
+            "datacomp/imagenet-1k-random-debug" + str(DEV_AMOUNT) + "-" + str(percentage), token=GATED_IMAGENET)
     else:
         Dataset.from_generator(updated_dataset.__iter__).push_to_hub(
             "datacomp/imagenet-1k-random" + str(percentage), token=GATED_IMAGENET)
@@ -63,6 +78,23 @@ def main(percentage=10):
     print("That took %d seconds" % (end - start))
 
 
+def get_label_fractions(dataset):
+    print("Getting label proportions.")
+    if READ_DISTRO:
+        with open("label_distro.json", "r+") as f:
+            label_counts = json.loads(f.read())
+    else:
+        label_counts = Counter([example['label'] for example in dataset])
+        # Don't overrwrite the distribution when devving.
+        if not DEV:
+            with open("label_distro.json", "w+") as f:
+                f.write(json.dumps(label_counts))
+    label_list = list(label_counts.keys())
+    denom = sum(label_counts.values())
+    label_fractions = [label_counts[key]/denom for key in label_keys]
+    return label_list, label_fractions
+
+
 def randomize_labels(examples, indices):
     # What set of examples should be randomized in this batch?
     # This is the intersection of the batch indices and the indices we randomly selected to change the labels of.
@@ -72,9 +104,7 @@ def randomize_labels(examples, indices):
         # Change the label to a random integer between 0 and 9
         for n in range(len(indices)):
             index = indices[n]
-            examples["label"][n] = random.randint(0,
-                                                  9) if index in batch_subset else \
-            examples["label"][n]
+            examples["label"][n] = new_random_labels.pop() if index in batch_subset else examples["label"][n]
     return examples
 
 demo = gr.Interface(fn=main, inputs="text", outputs="text")