Update app.py

app.py

@@ -24,6 +24,7 @@ https://huggingface.co/datasets/datacomp
 # https://huggingface.co/datasets/ILSVRC/imagenet-1k
 NUM_EXAMPLES = 1281167
 DEV = False
+FRACTIONS = [2, 4, 8, 16, 32, 64]
 # Arbitrary small number of dataset examples to look at, only using in devv'ing.
 DEV_AMOUNT = 10
 if DEV:
@@ -31,87 +32,38 @@ if DEV:
 # Whether to read in the distribution over labels from an external text file.
 READ_DISTRO = False
 GATED_IMAGENET = os.environ.get("GATED_IMAGENET")
-LABELS_FILE = "label_frequencies_full.csv"
-
-
-def read_label_frequencies():
-    label_counts_dict = {}
-    header_row = ['Label', 'Frequency']
-    with open(LABELS_FILE) as csvfile:
-        label_reader = csv.DictReader(csvfile)
-        assert label_reader.fieldnames == header_row
-        for row in label_reader:
-            assert row['Label'] not in label_counts_dict
-            label_counts_dict[row['Label']] = int(row['Frequency'])
-    # TODO: Can we just do this instead of the fractions? Do they really need to be normalized?      
-    # label_list, label_counts = zip(*label_counts_dict.items())
-    return label_counts_dict
-
-
-def get_label_fractions(label_counts_dict):
-    print("Getting label proportions.")
-    label_list = list(label_counts_dict.keys())
-    denom = sum(label_counts_dict.values())
-    label_fractions = [label_counts_dict[key]/denom for key in label_counts_dict]
-    return label_list, label_fractions
-
-
-def randomize_labels(examples, indices, new_random_labels):
-    # 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.
-    batch_subset = list(set(indices) & randomize_subset)
-    # If this batch has indices that we're changing the label of....
-    if batch_subset != []:
-        # Change the label to a random integer between 0 and 9
-        for n in range(len(indices)):
-            index = indices[n]
-            examples["label"][n] = new_random_labels.pop() if index in batch_subset else examples["label"][n]
-    return examples
 
 
+def create_subset_dataset(dataset, fraction_size):
+    dataset = dataset.shuffle(buffer_size=NUM_EXAMPLES)
+    num_samples = int(NUM_EXAMPLES / fraction_size)
+    sampled_dataset = dataset.take(num_samples)   
+    return sampled_dataset
+    
 def main(percentage=10):
     global randomize_subset
     # Just for timing how long this takes.
     start = time.time()
 
     percentage = float(percentage)
-    print("Randomizing %d percent of the data." % percentage)
-    # Set the random seed, based on the percentage, so that our random changes are reproducible.
-    random.seed(percentage)
-
-    # Load the dataset from the HF hub. Use streaming so as not to load the entire dataset at once.
-    # Use the .take(DEV_AMOUNT) to only grab a small chunk of instances to develop with.
+    
     if DEV:
-        dataset = load_dataset("ILSVRC/imagenet-1k", split="train", streaming=True,
+        dataset = load_dataset("datacomp/imagenet-1k-random" + str(percentage), split="train", streaming=True,
                            trust_remote_code=True, token=GATED_IMAGENET).take(DEV_AMOUNT)
     else:
-        dataset = load_dataset("ILSVRC/imagenet-1k", split="train", streaming=True,
+        dataset = load_dataset("datacomp/imagenet-1k-random" + str(percentage), split="train", streaming=True,
                            trust_remote_code=True, token=GATED_IMAGENET)
 
-    label_list, label_fractions = get_label_fractions(read_label_frequencies())
-
-    # How many new random labels are we creating?
-    num_new_labels = int(round(NUM_EXAMPLES * float(percentage) * .01))
-    
-    # Create a set of indices that are randomly chosen, to change their labels.
-    # 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_list, size=num_new_labels, p=label_fractions))
+    for frac in FRACTIONS:
+        sampled_dataset = create_subset_dataset(dataset, frac)
 
-    # Update the dataset so that the labels are randomized
-    updated_dataset = dataset.map(randomize_labels, with_indices=True, 
-                                  fn_kwargs={'new_random_labels':new_random_labels}, 
-                                  features=dataset.features, batched=True)
-
-    # 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(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)
+        # Upload the new version of the dataset (this will take awhile)
+        if DEV:
+            Dataset.from_generator(updated_dataset.__iter__).push_to_hub(
+                "datacomp/debug-imagenet-1k-random-" + "-" + str(percentage) + '-frac-1/' + str(frac), token=GATED_IMAGENET)
+        else:
+            Dataset.from_generator(updated_dataset.__iter__).push_to_hub(
+                "datacomp/imagenet-1k-random" + str(percentage) + '-frac-1/' + str(frac), token=GATED_IMAGENET)
 
 
     end = time.time()