Implement custom Huggingface dataset with data downloaded from s3

user_5296106 0
Viewed 392

In order to implement a custom Huggingface dataset I need to implement three methods:

from datasets import DatasetBuilder, DownloadManager

class MyDataset(DatasetBuilder):
    def _info(self):
        ...

    def _split_generator(self, dl_manager: DownloadManager):
        '''
        Method in charge of downloading (or retrieving locally
        the data files), organizing them according to the splits
        and defining specific arguments for the generation process
        if needed.
        '''
        ...

    def _generate_examples():
        ...

Now, in the _split_generator method I need to download a CSV file from S3 (a private bucket, one needs keys to access it). This file will be then further processed once it's been downloaded.

Do you know if there is a way to use the parameter dl_manager to download it? I guess I can download the file with some other methods/external libraries, but I'm wondering if I can do it with Huggingface's datasets objects and functionalities.

In this repo you can see many examples of custom datasets. For instance, the data used to build the amazon us reviews is downloaded from https://s3.amazonaws.com/amazon-reviews-pds/tsv/amazon_reviews_us_" + name + ".tsv.gz" (as you can see here). This is a public link though, and it can be accessed by everybody. Instead, I would like to use a Downloadmanager object to download my private data from S3.

huggingface-datasets amazon-s3 python
1 Answers

datasets provides some class to download things from S3 (and other Cloud providers) : https://huggingface.co/docs/datasets/v2.4.0/en/filesystems

So you can do something like :

def _split_generators(self, dl_manager):
    s3 = datasets.filesystems.S3FileSystem()

    _, f = os.path.split(MY_S3_URI)
    s3.get(MY_S3_URI, os.path.join(CACHE_DIR, f))

    return [
        datasets.SplitGenerator(name=datasets.Split.ALL, gen_kwargs={"filepath": os.path.join(CACHE_DIR, f)}),
    ]
user_94947900
Related