Process

Process

🌍 Datasets provides many tools for modifying the structure and content of a dataset. These tools are important for tidying up a dataset, creating additional columns, converting between features and formats, and much more.

This guide will show you how to:

  • Reorder rows and split the dataset.

  • Rename and remove columns, and other common column operations.

  • Apply processing functions to each example in a dataset.

  • Concatenate datasets.

  • Apply a custom formatting transform.

  • Save and export processed datasets.

For more details specific to processing other dataset modalities, take a look at the process audio dataset guide, the process image dataset guide, or the process text dataset guide.

The examples in this guide use the MRPC dataset, but feel free to load any dataset of your choice and follow along!

Copied

>>> from datasets import load_dataset
>>> dataset = load_dataset("glue", "mrpc", split="train")

All processing methods in this guide return a new Dataset object. Modification is not done in-place. Be careful about overriding your previous dataset!

Sort, shuffle, select, split, and shard

There are several functions for rearranging the structure of a dataset. These functions are useful for selecting only the rows you want, creating train and test splits, and sharding very large datasets into smaller chunks.

Sort

Use sort() to sort column values according to their numerical values. The provided column must be NumPy compatible.

Copied

Under the hood, this creates a list of indices that is sorted according to values of the column. This indices mapping is then used to access the right rows in the underlying Arrow table.

Shuffle

The shuffle() function randomly rearranges the column values. You can specify the generator parameter in this function to use a different numpy.random.Generator if you want more control over the algorithm used to shuffle the dataset.

Copied

Shuffling takes the list of indices [0:len(my_dataset)] and shuffles it to create an indices mapping. However as soon as your Dataset has an indices mapping, the speed can become 10x slower. This is because there is an extra step to get the row index to read using the indices mapping, and most importantly, you aren’t reading contiguous chunks of data anymore. To restore the speed, you’d need to rewrite the entire dataset on your disk again using Dataset.flatten_indices(), which removes the indices mapping. Alternatively, you can switch to an IterableDataset and leverage its fast approximate shuffling IterableDataset.shuffle():

Copied

Select and Filter

There are two options for filtering rows in a dataset: select() and filter().

  • select() returns rows according to a list of indices:

Copied

  • filter() returns rows that match a specified condition:

Copied

filter() can also filter by indices if you set with_indices=True:

Copied

Unless the list of indices to keep is contiguous, those methods also create an indices mapping under the hood.

Split

The train_test_split() function creates train and test splits if your dataset doesn’t already have them. This allows you to adjust the relative proportions or an absolute number of samples in each split. In the example below, use the test_size parameter to create a test split that is 10% of the original dataset:

Copied

The splits are shuffled by default, but you can set shuffle=False to prevent shuffling.

Shard

🌍 Datasets supports sharding to divide a very large dataset into a predefined number of chunks. Specify the num_shards parameter in shard() to determine the number of shards to split the dataset into. You’ll also need to provide the shard you want to return with the index parameter.

For example, the imdb dataset has 25000 examples:

Copied

After sharding the dataset into four chunks, the first shard will only have 6250 examples:

Copied

Rename, remove, cast, and flatten

The following functions allow you to modify the columns of a dataset. These functions are useful for renaming or removing columns, changing columns to a new set of features, and flattening nested column structures.

Rename

Use rename_column() when you need to rename a column in your dataset. Features associated with the original column are actually moved under the new column name, instead of just replacing the original column in-place.

Provide rename_column() with the name of the original column, and the new column name:

Copied

Remove

When you need to remove one or more columns, provide the column name to remove to the remove_columns() function. Remove more than one column by providing a list of column names:

Copied

Conversely, select_columns() selects one or more columns to keep and removes the rest. This function takes either one or a list of column names:

Copied

Cast

The cast() function transforms the feature type of one or more columns. This function accepts your new Features as its argument. The example below demonstrates how to change the ClassLabel and Value features:

Copied

Casting only works if the original feature type and new feature type are compatible. For example, you can cast a column with the feature type Value("int32") to Value("bool") if the original column only contains ones and zeros.

Use the cast_column() function to change the feature type of a single column. Pass the column name and its new feature type as arguments:

Copied

Flatten

Sometimes a column can be a nested structure of several types. Take a look at the nested structure below from the SQuAD dataset:

Copied

The answers field contains two subfields: text and answer_start. Use the flatten() function to extract the subfields into their own separate columns:

Copied

Notice how the subfields are now their own independent columns: answers.text and answers.answer_start.

Map

Some of the more powerful applications of 🌍 Datasets come from using the map() function. The primary purpose of map() is to speed up processing functions. It allows you to apply a processing function to each example in a dataset, independently or in batches. This function can even create new rows and columns.

In the following example, prefix each sentence1 value in the dataset with 'My sentence: '.

Start by creating a function that adds 'My sentence: ' to the beginning of each sentence. The function needs to accept and output a dict:

Copied

Now use map() to apply the add_prefix function to the entire dataset:

Copied

Let’s take a look at another example, except this time, you’ll remove a column with map(). When you remove a column, it is only removed after the example has been provided to the mapped function. This allows the mapped function to use the content of the columns before they are removed.

Specify the column to remove with the remove_columns parameter in map():

Copied

🌍 Datasets also has a remove_columns() function which is faster because it doesn’t copy the data of the remaining columns.

You can also use map() with indices if you set with_indices=True. The example below adds the index to the beginning of each sentence:

Copied

The map() also works with the rank of the process if you set with_rank=True. This is analogous to the with_indices parameter. The with_rank parameter in the mapped function goes after the index one if it is already present.

Copied

The main use-case for rank is to parallelize computation across several GPUs. This requires setting multiprocess.set_start_method("spawn"). If you don’t you’ll receive the following CUDA error:

Copied

Multiprocessing

Multiprocessing significantly speeds up processing by parallelizing processes on the CPU. Set the num_proc parameter in map() to set the number of processes to use:

Copied

Batch processing

The map() function supports working with batches of examples. Operate on batches by setting batched=True. The default batch size is 1000, but you can adjust it with the batch_size parameter. Batch processing enables interesting applications such as splitting long sentences into shorter chunks and data augmentation.

Split long examples

When examples are too long, you may want to split them into several smaller chunks. Begin by creating a function that:

  1. Splits the sentence1 field into chunks of 50 characters.

  2. Stacks all the chunks together to create the new dataset.

Copied

Apply the function with map():

Copied

Notice how the sentences are split into shorter chunks now, and there are more rows in the dataset.

Copied

Data augmentation

The map() function could also be used for data augmentation. The following example generates additional words for a masked token in a sentence.

Load and use the RoBERTA model in 🌍 Transformers’ FillMaskPipeline:

Copied

Create a function to randomly select a word to mask in the sentence. The function should also return the original sentence and the top two replacements generated by RoBERTA.

Copied

Use map() to apply the function over the whole dataset:

Copied

For each original sentence, RoBERTA augmented a random word with three alternatives. The original word distorting is supplemented by withholding, suppressing, and destroying.

Process multiple splits

Many datasets have splits that can be processed simultaneously with DatasetDict.map(). For example, tokenize the sentence1 field in the train and test split by:

Copied

Distributed usage

When you use map() in a distributed setting, you should also use torch.distributed.barrier. This ensures the main process performs the mapping, while the other processes load the results, thereby avoiding duplicate work.

The following example shows how you can use torch.distributed.barrier to synchronize the processes:

Copied

Concatenate

Separate datasets can be concatenated if they share the same column types. Concatenate datasets with concatenate_datasets():

Copied

You can also concatenate two datasets horizontally by setting axis=1 as long as the datasets have the same number of rows:

Copied

Interleave

You can also mix several datasets together by taking alternating examples from each one to create a new dataset. This is known as interleaving, which is enabled by the interleave_datasets() function. Both interleave_datasets() and concatenate_datasets() work with regular Dataset and IterableDataset objects. Refer to the Stream guide for an example of how to interleave IterableDataset objects.

You can define sampling probabilities for each of the original datasets to specify how to interleave the datasets. In this case, the new dataset is constructed by getting examples one by one from a random dataset until one of the datasets runs out of samples.

Copied

You can also specify the stopping_strategy. The default strategy, first_exhausted, is a subsampling strategy, i.e the dataset construction is stopped as soon one of the dataset runs out of samples. You can specify stopping_strategy=all_exhausted to execute an oversampling strategy. In this case, the dataset construction is stopped as soon as every samples in every dataset has been added at least once. In practice, it means that if a dataset is exhausted, it will return to the beginning of this dataset until the stop criterion has been reached. Note that if no sampling probabilities are specified, the new dataset will have max_length_datasets*nb_dataset samples.

Copied

Format

The set_format() function changes the format of a column to be compatible with some common data formats. Specify the output you’d like in the type parameter and the columns you want to format. Formatting is applied on-the-fly.

For example, create PyTorch tensors by setting type="torch":

Copied

The with_format() function also changes the format of a column, except it returns a new Dataset object:

Copied

🌍 Datasets also provides support for other common data formats such as NumPy, Pandas, and JAX. Check out the Using Datasets with TensorFlow guide for more details on how to efficiently create a TensorFlow dataset.

If you need to reset the dataset to its original format, use the reset_format() function:

Copied

Format transform

The set_transform() function applies a custom formatting transform on-the-fly. This function replaces any previously specified format. For example, you can use this function to tokenize and pad tokens on-the-fly. Tokenization is only applied when examples are accessed:

Copied

You can also use the set_transform() function to decode formats not supported by Features. For example, the Audio feature uses soundfile - a fast and simple library to install - but it does not provide support for less common audio formats. Here is where you can use set_transform() to apply a custom decoding transform on the fly. You’re free to use any library you like to decode the audio files.

The example below uses the pydub package to open an audio format not supported by soundfile:

Copied

Save

Once you are done processing your dataset, you can save and reuse it later with save_to_disk().

Save your dataset by providing the path to the directory you wish to save it to:

Copied

Use the load_from_disk() function to reload the dataset:

Copied

Want to save your dataset to a cloud storage provider? Read our Cloud Storage guide to learn how to save your dataset to AWS or Google Cloud Storage.

Export

🌍 Datasets supports exporting as well so you can work with your dataset in other applications. The following table shows currently supported file formats you can export to:

File type
Export method

CSV

Dataset.to_csv()

JSON

Dataset.to_json()

Parquet

Dataset.to_parquet()

SQL

Dataset.to_sql()

In-memory Python object

Dataset.to_pandas() or Dataset.to_dict()

For example, export your dataset to a CSV file like this:

Copied

PreviousLoadNextStream

Last updated