Using the Hugging Face Datasets library with PyTorch

Using the Hugging Face datasets library with PyTorch enables easy access to thousands of public datasets and simplifies handling custom ones. There are over 144,000 (in May 2024) datasets available on Hugging Face, which can be checked with the following lines of code:

from huggingface_hub import hf_api
datasets = hf_api.list_datasets()
len([d for d in datasets])

To get started with the Hugging Face datasets library, make sure you have installed the following dependencies:

pip install torch datasets transformers

All code for this section is available on GitHub [9]. First, we should import the required libraries and set up the environment:

import torch
from datasets import load_dataset
from transformers import BertTokenizer

We import the load_dataset function from the datasets library. We plan to use the BERT model for our demonstration, hence we import the BertTokenizer to convert text into tokens.

...