Données massives ? 🤗 <i> Datasets </i> à la rescousse !

De nos jours, il n’est pas rare de travailler avec des jeux de données de plusieurs gigaoctets surtout si vous envisagez de pré-entraîner un transformer comme BERT ou GPT-2 à partir de zéro. Dans ces cas, même charger les données peut être un défi. Par exemple, le corpus WebText utilisé pour pré-entraîner GPT-2 se compose de plus de 8 millions de documents et de 40 Go de texte. Le charger dans la RAM de votre ordinateur portable est susceptible de lui donner une crise cardiaque !

Heureusement, 🤗 Datasets a été conçu pour surmonter ces limitations. Il vous libère des problèmes de gestion de la mémoire en traitant les jeux de données comme des fichiers mappés en mémoire, ainsi que des limites du disque dur en faisant du streaming sur les entrées dans un corpus.

Dans cette section, nous allons explorer ces fonctionnalités de 🤗 Datasets avec un énorme corpus de 825 Go connu sous le nom de The Pile. Commençons !

Qu’est-ce que <i> The Pile </i> ?

The Pile est un corpus de texte en anglais créé par EleutherAI pour entraîner des modèles de langage à grande échelle. Il comprend une gamme variée de jeux de données, couvrant des articles scientifiques, des référentiels de code GitHub et du texte Web filtré. Le corpus d’entraînement est disponible en morceaux de 14 Go et vous pouvez aussi télécharger plusieurs des composants individuels. Commençons par jeter un coup d’œil au jeu de données PubMed Abstracts, qui est un corpus de résumés de 15 millions de publications biomédicales sur PubMed. Le jeu de données est au format JSON Lines et est compressé à l’aide de la bibliothèque zstandard. Nous devons donc d’abord installer cette bibliothèque :

!pip install zstandard

Ensuite, nous pouvons charger le jeu de données en utilisant la méthode pour les fichiers distants que nous avons apprise dans section 2 :

from datasets import load_dataset

# Cela prend quelques minutes à exécuter, alors allez prendre un thé ou un café en attendant :)
data_files = "https://the-eye.eu/public/AI/pile_preliminary_components/PUBMED_title_abstracts_2019_baseline.jsonl.zst"
pubmed_dataset = load_dataset("json", data_files=data_files, split="train")
pubmed_dataset

Dataset({
    features: ['meta', 'text'],
    num_rows: 15518009
})

Nous pouvons voir qu’il y a 15 518 009 lignes et 2 colonnes dans notre jeu de données. C’est beaucoup !

Inspectons le contenu du premier exemple :

pubmed_dataset[0]

{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'
# Épidémiologie de l'hypoxémie chez les enfants souffrant d'une infection aiguë des voies respiratoires inférieures. Déterminer la prévalence de l'hypoxémie chez les enfants de moins de 5 ans souffrant d'une infection aiguë des voies respiratoires inférieures (IAVI), les facteurs de risque de l'hypoxémie chez les enfants de moins de 5 ans souffrant d'une IAVI, et l'association de l'hypoxémie à un risque accru de décès chez les enfants du même âge ...
}

Cela ressemble au résumé d’un article médical. Voyons maintenant combien de RAM nous avons utilisé pour charger le jeu de données !

La magie du <i> memory mapping </i>

Un moyen simple de mesurer l’utilisation de la mémoire dans Python consiste à utiliser la bibliothèque psutil qui peut être installée avec pip comme suit :

!pip install psutil

Elle fournit une classe Process qui permet de vérifier l’utilisation de la mémoire du processus en cours :

import psutil

# Process.memory_info est exprimé en octets, donc convertir en mégaoctets
print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

RAM used: 5678.33 MB

Ici, l’attribut rss fait référence à la taille de l’ensemble résident, qui est la fraction de mémoire qu’un processus occupe dans la RAM. Cette mesure inclut également la mémoire utilisée par l’interpréteur Python et les bibliothèques que nous avons chargées, de sorte que la quantité réelle de mémoire utilisée pour charger le jeu de données est un peu plus petite. À titre de comparaison, voyons la taille du jeu de données sur le disque en utilisant l’attribut dataset_size. Comme le résultat est exprimé en octets comme précédemment, nous devons le convertir manuellement en gigaoctets :

print(f"Number of files in dataset : {pubmed_dataset.dataset_size}")
size_gb = pubmed_dataset.dataset_size / (1024**3)
print(f"Dataset size (cache file) : {size_gb:.2f} GB")

Number of files in dataset : 20979437051
Dataset size (cache file) : 19.54 GB

Malgré sa taille de près de 20 Go, nous pouvons charger et accéder au jeu de données avec beaucoup moins de RAM !

Si vous êtes familier avec Pandas, ce résultat pourrait surprendre en raison de la célèbre règle d’or de Wes Kinney selon laquelle vous avez généralement besoin de 5 à 10 fois plus de RAM que la taille de votre jeu de données. Alors, comment 🤗 Datasets résout-il ce problème de gestion de la mémoire ? 🤗 Datasets traite chaque jeu de données comme un fichier mappé en mémoire. Cela fournit un mappage entre la RAM et le stockage du système de fichiers permettant à la bibliothèque d’accéder et d’opérer sur des éléments du jeu de données sans avoir besoin de le charger entièrement en mémoire.

Les fichiers mappés en mémoire peuvent également être partagés entre plusieurs processus ce qui permet de paralléliser des méthodes telles que Dataset.map() sans avoir à déplacer ou copier le jeu de données. Sous le capot, ces capacités sont toutes réalisées par le format de mémoire Apache Arrow et pyarrow, qui accélèrent le chargement et le traitement des données. (Pour plus de détails sur Apache Arrow et les comparaisons avec Pandas, consultez l’article de blog de Dejan Simic). Pour voir ceci en action, effectuons un petit test de vitesse en itérant sur tous les éléments du jeu de données PubMed Abstracts :

import timeit

code_snippet = """batch_size = 1000

for idx in range(0, len(pubmed_dataset), batch_size):
    _ = pubmed_dataset[idx:idx + batch_size]
"""

time = timeit.timeit(stmt=code_snippet, number=1, globals=globals())
print(
    f"Iterated over {len(pubmed_dataset)} examples (about {size_gb:.1f} GB) in "
    f"{time:.1f}s, i.e. {size_gb/time:.3f} GB/s"
)

'Iterated over 15518009 examples (about 19.5 GB) in 64.2s, i.e. 0.304 GB/s'

Ici, nous avons utilisé le module timeit de Python pour mesurer le temps d’exécution pris par code_snippet. Vous pourrez généralement itérer sur un jeu de données à une vitesse de quelques dixièmes de Go/s à plusieurs Go/s. Cela fonctionne très bien pour la grande majorité des applications, mais vous devrez parfois travailler avec un jeu de données trop volumineux pour être même stocké sur le disque dur de votre ordinateur portable. Par exemple, si nous essayions de télécharger The Pile dans son intégralité, nous aurions besoin de 825 Go d’espace disque libre ! Pour gérer ces cas, 🤗 Datasets fournit une fonctionnalité de streaming qui nous permet de télécharger et d’accéder aux éléments à la volée, sans avoir besoin de télécharger l’intégralité du jeu de données. Voyons comment cela fonctionne.

Jeux de données en continu

Pour activer le streaming du jeu de données, il vous suffit de passer l’argument streaming=True à la fonction load_dataset(). Par exemple, chargeons à nouveau le jeu de données PubMed Abstracts mais en mode streaming :

pubmed_dataset_streamed = load_dataset(
    "json", data_files=data_files, split="train", streaming=True
)

Au lieu du familier Dataset que nous avons rencontré ailleurs dans ce chapitre, l’objet retourné avec streaming=True est un IterableDataset. Comme son nom l’indique, pour accéder aux éléments d’un IterableDataset, nous devons parcourir celui-ci. Nous pouvons accéder au premier élément de notre jeu de données diffusé comme suit :

next(iter(pubmed_dataset_streamed))

{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

Les éléments d’un jeu de données diffusé en continu peuvent être traités à la volée à l’aide de IterableDataset.map(), ce qui est utile pendant l’entraînement si vous avez besoin de tokeniser les entrées. Le processus est exactement le même que celui que nous avons utilisé pour tokeniser notre jeu de données dans le chapitre 3, à la seule différence que les sorties sont renvoyées une par une :

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
tokenized_dataset = pubmed_dataset_streamed.map(lambda x: tokenizer(x["text"]))
next(iter(tokenized_dataset))

{'input_ids': [101, 4958, 5178, 4328, 6779, ...], 'attention_mask': [1, 1, 1, 1, 1, ...]}

Vous pouvez également mélanger un jeu de données diffusé en continu à l’aide de IterableDataset.shuffle(), mais contrairement à Dataset.shuffle(), cela ne mélange que les éléments dans un buffer_size prédéfini :

shuffled_dataset = pubmed_dataset_streamed.shuffle(buffer_size=10_000, seed=42)
next(iter(shuffled_dataset))

{'meta': {'pmid': 11410799, 'language': 'eng'},
 'text': 'Randomized study of dose or schedule modification of granulocyte colony-stimulating factor in platinum-based chemotherapy for elderly patients with lung cancer ...'
# Étude randomisée sur la modification de la dose ou du calendrier d'administration du facteur de stimulation des colonies de granulocytes dans le cadre d'une chimiothérapie à base de platine chez les patients âgés atteints de cancer du poumon ...
}

Dans cet exemple, nous avons sélectionné un exemple aléatoire parmi les 10 000 premiers exemples du tampon. Une fois qu’un exemple est accédé, sa place dans le tampon est remplie avec l’exemple suivant dans le corpus (c’est-à-dire le 10 001e exemple dans le cas ci-dessus). Vous pouvez également sélectionner des éléments d’un jeu de données diffusé en continu à l’aide des fonctions IterableDataset.take() et IterableDataset.skip(), qui agissent de la même manière que Dataset.select(). Par exemple, pour sélectionner les 5 premiers exemples dans le jeu de données PubMed Abstracts, nous pouvons procéder comme suit :

dataset_head = pubmed_dataset_streamed.take(5)
list(dataset_head)

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'
# Épidémiologie de l'hypoxémie chez les enfants atteints d'une infection aiguë des voies respiratoires inférieures ...},
 {'meta': {'pmid': 11409575, 'language': 'eng'},
  'text': 'Clinical signs of hypoxaemia in children with acute lower respiratory infection: indicators of oxygen therapy ...'
# Signes cliniques d'hypoxémie chez les enfants atteints d'une infection aiguë des voies respiratoires inférieures : indicateurs de l'oxygénothérapie ...},
 {'meta': {'pmid': 11409576, 'language': 'eng'},
  'text': "Hypoxaemia in children with severe pneumonia in Papua New Guinea ..."
# Hypoxémie chez les enfants atteints de pneumonie grave en Papouasie-Nouvelle-Guinée ...},
 {'meta': {'pmid': 11409577, 'language': 'eng'},
  'text': 'Oxygen concentrators and cylinders ...'
# Concentrateurs et bouteilles d'oxygène...},
 {'meta': {'pmid': 11409578, 'language': 'eng'},
  'text': 'Oxygen supply in rural africa: a personal experience ...'
# L'approvisionnement en oxygène dans les zones rurales africaines : une expérience personnelle ...}]

De même, vous pouvez utiliser la fonction IterableDataset.skip() pour créer des fractionnements d’entraînement et de validation à partir d’un jeu de données mélangé comme suit :

# Ignorer les 1 000 premiers exemples et inclure le reste dans l'ensemble d'apprentissage.
train_dataset = shuffled_dataset.skip(1000)
# Prendre les 1 000 premiers exemples pour l'ensemble de validation.
validation_dataset = shuffled_dataset.take(1000)

Terminons notre exploration du streaming des jeux de données avec une application commune : combiner plusieurs jeux de données pour créer un seul corpus. 🤗 Datasets fournit une fonction interleave_datasets() qui convertit une liste d’objets IterableDataset en un seul IterableDataset, où les éléments du nouveau jeu de données sont obtenus en alternant entre les exemples source. Cette fonction est particulièrement utile lorsque vous essayez de combiner de grands jeux de données. Par exemple, streamons FreeLaw, un sous-ensemble de The Pile et qui est un jeu de données de 51 Go d’avis juridiques de tribunaux américains :

law_dataset_streamed = load_dataset(
    "json",
    data_files="https://the-eye.eu/public/AI/pile_preliminary_components/FreeLaw_Opinions.jsonl.zst",
    split="train",
    streaming=True,
)
next(iter(law_dataset_streamed))

{'meta': {'case_ID': '110921.json',
  'case_jurisdiction': 'scotus.tar.gz',
  'date_created': '2010-04-28T17:12:49Z'},
 'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}

Ce jeu de données est suffisamment volumineux pour solliciter la RAM de la plupart des ordinateurs portables, mais nous avons pu le charger et y accéder sans transpirer ! Combinons maintenant les jeux de données FreeLaw et PubMed Abstracts avec la fonction interleave_datasets() :

from itertools import islice
from datasets import interleave_datasets

combined_dataset = interleave_datasets([pubmed_dataset_streamed, law_dataset_streamed])
list(islice(combined_dataset, 2))

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'case_ID': '110921.json',
   'case_jurisdiction': 'scotus.tar.gz',
   'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}]

Ici, nous avons utilisé la fonction islice() du module itertools de Python pour sélectionner les deux premiers exemples du jeu de données combiné. Nous pouvons voir qu’ils correspondent aux premiers exemples de chacun des deux jeux de données source.

Enfin, si vous souhaitez streamer The Pile dans son intégralité de 825 Go, vous pouvez récupérer tous les fichiers préparés comme suit :

base_url = "https://the-eye.eu/public/AI/pile/"
data_files = {
    "train": [base_url + "train/" + f"{idx:02d}.jsonl.zst" for idx in range(30)],
    "validation": base_url + "val.jsonl.zst",
    "test": base_url + "test.jsonl.zst",
}
pile_dataset = load_dataset("json", data_files=data_files, streaming=True)
next(iter(pile_dataset["train"]))

{'meta': {'pile_set_name': 'Pile-CC'},
 'text': 'It is done, and submitted. You can play “Survival of the Tastiest” on Android, and on the web...'}

Vous disposez maintenant de tous les outils dont vous avez besoin pour charger et traiter des jeux de données de toutes formes et tailles. Cependant à moins que vous ne soyez exceptionnellement chanceux, il arrivera un moment dans votre cheminement en traitement du langage naturel où vous devrez réellement créer un jeu de données pour résoudre un problème donné. C’est le sujet de la section suivante !