datacomp/select-subset

1import os2import json3import csv4 5import gradio as gr6import random7import time8 9from collections import Counter10from numpy.random import choice11from datasets import load_dataset, Dataset12 13from PIL import PngImagePlugin, ImageFile14ImageFile.LOAD_TRUNCATED_IMAGES = True15PngImagePlugin.MAX_TEXT_CHUNK = 1048576 * 10 # this is 10x the amount.16"""17This code is designed to read in the ImageNet 1K ILSVRC dataset from the Hugging Face Hub, 18then create a new version of this dataset with {percentage} lines with random labels based on the observed frequencies,19then upload this new version of the Hugging Face Hub, in the Data Composition organization:20https://huggingface.co/datasets/datacomp21"""22 23# The number of examples/instances in this dataset is copied from the model card:24# https://huggingface.co/datasets/ILSVRC/imagenet-1k25NUM_EXAMPLES = 128116726DEV = False27FRACTIONS = [2, 4, 8, 16, 32, 64]28# Arbitrary small number of dataset examples to look at, only using in devv'ing.29DEV_AMOUNT = 1030if DEV:31    NUM_EXAMPLES = DEV_AMOUNT32    FRACTIONS = [2]33# Whether to read in the distribution over labels from an external text file.34READ_DISTRO = False35GATED_IMAGENET = os.environ.get("GATED_IMAGENET")36 37 38def create_subset_dataset(dataset, fraction_size):39    dataset = dataset.shuffle()40    num_samples = int(NUM_EXAMPLES / fraction_size)41    sampled_dataset = dataset.take(num_samples)   42    return sampled_dataset43    44def main(percentage=10):45    global randomize_subset46    # Just for timing how long this takes.47    start = time.time()48 49    percentage = float(percentage)50    51    if DEV:52        dataset = load_dataset("datacomp/imagenet-1k-random" + str(percentage), split="train", streaming=True,53                           trust_remote_code=True, token=GATED_IMAGENET).take(DEV_AMOUNT)54    else:55        dataset = load_dataset("datacomp/imagenet-1k-random" + str(percentage) + "-take2", split="train", streaming=True,56                           trust_remote_code=True, token=GATED_IMAGENET)57 58    for frac in FRACTIONS:59        sampled_dataset = create_subset_dataset(dataset, frac)60 61        # Upload the new version of the dataset (this will take awhile)62        if DEV:63            Dataset.from_generator(sampled_dataset.__iter__).push_to_hub(64                "datacomp/debug-imagenet-1k-random-" + str(percentage) + '-frac-1over' + str(frac), token=GATED_IMAGENET)65        else:66            Dataset.from_generator(sampled_dataset.__iter__).push_to_hub(67                "datacomp/imagenet-1k-random-" + str(percentage) + '-frac-1over' + str(frac) + "-take2", token=GATED_IMAGENET)68 69 70    end = time.time()71    print("That took %d seconds" % (end - start))72 73 74demo = gr.Interface(fn=main, inputs="text", outputs="text")75demo.launch()76