Datasets documentation
Stream
Stream
Le streaming de jeu de données vous permet de travailler avec un jeu de données sans le télécharger. Les données sont diffusées au fur et à mesure que vous itérez sur le jeu de données. Cela est particulièrement utile lorsque :
- Vous ne souhaitez pas attendre le téléchargement d’un jeu de données extrêmement volumineux.
- La taille du jeu de données dépasse l'espace disque disponible sur votre ordinateur.
- Vous souhaitez explorer rapidement quelques échantillons d'un jeu de données.

Par exemple, la partie anglaise du jeu de données HuggingFaceFW/fineweb fait 45 téraoctets, mais vous pouvez l’utiliser instantanément avec le streaming. Pour streamer un jeu de données, définissez streaming=True dans load_dataset() comme indiqué ci-dessous :
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> print(next(iter(dataset)))
{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012\nWhen the first serious bout of tornadoes of 2012 blew through middle America in the middle of the night, they touched down in places hours from any AP bureau...', ...,
'language_score': 0.9721424579620361, 'token_count': 717}Le streaming de jeux de données vous permet également de travailler avec un jeu de données composé de fichiers locaux sans effectuer aucune conversion. Dans ce cas, les données sont diffusées depuis les fichiers locaux au fur et à mesure que vous itérez sur le jeu de données. Cela est particulièrement utile lorsque :
- Vous ne souhaitez pas attendre qu’un très grand jeu de données local soit converti au format Arrow.
- La taille des fichiers convertis dépasserait l'espace disque disponible sur votre ordinateur.
- Vous souhaitez explorer rapidement quelques échantillons d'un jeu de données.
- Vous souhaitez charger uniquement certaines colonnes ou filtrer efficacement un jeu de données Parquet.
Par exemple, vous pouvez diffuser un jeu de données local composé de centaines de fichiers JSONL compressés, comme oscar-corpus/OSCAR-2201, pour l’utiliser immédiatement :
>>> from datasets import load_dataset
>>> data_files = {'train': 'path/to/OSCAR-2201/compressed/en_meta/*.jsonl.gz'}
>>> dataset = load_dataset('json', data_files=data_files, split='train', streaming=True)
>>> print(next(iter(dataset)))
{'id': 0, 'text': 'Founded in 2015, Golden Bees is a leading programmatic recruitment platform dedicated to employers, HR agencies and job boards. The company has developed unique HR-custom technologies and predictive algorithms to identify and attract the best candidates for a job opportunity.', ...Parquet est un format en colonnes qui permet de diffuser et de charger uniquement un sous-ensemble de colonnes et d’ignorer les colonnes non souhaitées. Parquet stocke également des métadonnées telles que les statistiques de colonnes (au niveau du fichier et du groupe de lignes), permettant un filtrage efficace. Utilisez les arguments columns et filters de datasets.packaged_modules.parquet.ParquetConfig pour diffuser des jeux de données Parquet, sélectionner des colonnes et appliquer des filtres :
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, columns=["url", "date"])
>>> print(next(iter(dataset)))
{'url': 'http://%[email protected]/Content/Press-Release/2012/How-AP-reported-in-all-formats-from-tornado-stricken-regions', 'date': '2013-05-18T05:48:54Z'}
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, filters=[("language_score", ">=", 0.99)])
>>> print(next(iter(dataset)))
{'text': 'Everyone wishes for something. And lots of people believe they know how to make their wishes come true with magical thinking.\nWhat is it? "Magical thinking is a belief in forms of causation, with no known physical basis," said Professor Emily Pronin of Princeton...', ...,
'language_score': 0.9900368452072144, 'token_count': 716}Charger un jeu de données en mode streaming crée une nouvelle instance de type de jeu de données (au lieu de l'objet classique Dataset), connue sous le nom de IterableDataset. Ce type de jeu de données particulier dispose de sa propre série de méthodes de traitement présentées ci-dessous.
Un IterableDataset est utile pour les tâches itératives comme l’entraînement d’un modèle. Vous ne devriez pas utiliser un IterableDataset pour les tâches qui nécessitent un accès aléatoire aux exemples, car vous devez parcourir l’ensemble de l’itérable à l’aide d’une boucle for. Obtenir le dernier exemple d’un jeu de données itérable vous obligerait à parcourir tous les exemples précédents. Vous trouverez plus de détails dans le guide Dataset vs. IterableDataset.
Indexation des colonnes
Il est parfois pratique d'itérer sur les valeurs d'une colonne spécifique. Heureusement, un IterableDataset prend en charge l'indexation par colonne :
>>> from datasets import load_dataset
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> print(next(iter(dataset["text"])))
Beginners BBQ Class Taking Place in Missoula!...Convertir depuis un Dataset
Si vous disposez déjà d'un objet Dataset, vous pouvez le convertir en IterableDataset à l'aide de la fonction to_iterable_dataset(). C'est en réalité plus rapide que de définir l'argument streaming=True dans load_dataset(), car les données sont diffusées depuis des fichiers locaux.
>>> from datasets import load_dataset
# faster 🐇
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset()
# slower 🐢
>>> iterable_dataset = load_dataset("ethz/food101", streaming=True)La fonction to_iterable_dataset() prend en charge le sharding lors de l’instanciation de l’IterableDataset. C’est utile lorsque vous travaillez avec de grands jeux de données et que vous souhaitez mélanger le jeu de données ou activer un chargement parallèle rapide avec un PyTorch DataLoader.
>>> import torch
>>> from datasets import load_dataset
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset(num_shards=64) # shard the dataset
>>> iterable_dataset = iterable_dataset.shuffle(buffer_size=10_000) # shuffles the shards order and use a shuffle buffer when you start iterating
dataloader = torch.utils.data.DataLoader(iterable_dataset, num_workers=4) # assigns 64 / 4 = 16 shards from the shuffled list of shards to each worker when you start iteratingMélanger
Comme pour un objet Dataset classique, vous pouvez également mélanger un IterableDataset avec IterableDataset.shuffle().
L’argument buffer_size contrôle la taille du tampon à partir duquel les exemples sont échantillonnés aléatoirement. Supposons que votre jeu de données contient un million d’exemples et que vous définissez buffer_size à dix mille. IterableDataset.shuffle() sélectionnera aléatoirement des exemples parmi les dix mille premiers exemples du tampon. Les exemples sélectionnés dans le tampon sont remplacés par de nouveaux exemples. Par défaut, la taille du tampon est de 1 000.
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> shuffled_dataset = dataset.shuffle(seed=42, buffer_size=10_000)IterableDataset.shuffle() remplit le tampon en utilisant jusqu'à
max_buffer_input_shardsshards à la fois (la valeur par défaut est 10), ce qui améliore considérablement la qualité du mélange. Il mélange également l'ordre des shards si le jeu de données est réparti sur plusieurs fichiers.
Si vous diffusez un jeu de données Parquet ne contenant qu’un seul fichier ou très peu de fichiers, utilisez reshard() avant le mélange. Cela permet de fractionner le jeu de données par groupe de lignes plutôt que par fichier, ce qui améliore la qualité du mélange.
Mélanger
Parfois, vous pouvez vouloir réorganiser le jeu de données après chaque époque. Cela nécessite de définir une graine différente pour chaque époque. Utilisez IterableDataset.set_epoch() entre les époques pour indiquer au jeu de données l'époque en cours.
Votre graine devient effectivement : initial seed + current epoch.
>>> for epoch in range(epochs):
... shuffled_dataset.set_epoch(epoch)
... for example in shuffled_dataset:
... ...Diviser le jeu de données
Vous pouvez diviser votre jeu de données de l'une des deux manières suivantes :
- IterableDataset.take() renvoie les
npremiers exemples d’un jeu de données :
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> dataset_head = dataset.take(2)
>>> list(dataset_head)
[{'text': "How AP reported in all formats from tor...},
{'text': 'Did you know you have two little yellow...}]- IterableDataset.skip() omet les
npremiers exemples d’un jeu de données et renvoie les exemples restants :
>>> train_dataset = shuffled_dataset.skip(1000)
takeetskipempêchent les appels ultérieurs àshufflecar ils verrouillent l’ordre des shards. Vous devezshufflevotre jeu de données avant de le diviser.
Shard
🤗 Datasets prend en charge le sharding pour diviser un très grand jeu de données en un nombre prédéfini de morceaux. Spécifiez le paramètre num_shards dans shard() pour déterminer le nombre de shards en lesquels diviser le jeu de données. Vous devez également fournir le shard que vous souhaitez renvoyer avec le paramètre index.
Par exemple, le jeu de données amazon_polarity comporte 4 shards (ici, il s'agit de 4 fichiers Parquet) :
>>> from datasets import load_dataset
>>> dataset = load_dataset("fancyzhx/amazon_polarity", split="train", streaming=True)
>>> print(dataset)
IterableDataset({
features: ['label', 'title', 'content'],
num_shards: 4
})Après avoir découpé le jeu de données en deux morceaux, le premier ne contiendra que 2 shards :
>>> dataset.shard(num_shards=2, index=0)
IterableDataset({
features: ['label', 'title', 'content'],
num_shards: 2
})Pour augmenter le nombre de shards d'un jeu de données, vous pouvez utiliser IterableDataset.reshard() :
>>> dataset.reshard()
IterableDataset({
features: ['label', 'title', 'content'],
num_shards: 3600
})Le mécanisme de re-sharding dépend du format des fichiers de jeu de données. Par exemple, pour Parquet, il effectue le re-sharding à l'aide de groupes de lignes au lieu d'avoir un fichier par shard. Consultez la documentation de IterableDataset.reshard() pour voir comment cela fonctionne pour chaque format.
Si votre jeu de données a dataset.num_shards==1 même après le re-sharding, vous devez le découper à l'aide de IterableDataset.skip() et IterableDataset.take() à la place.
Concaténer
Des jeux de données distincts peuvent être concaténés s'ils partagent les mêmes types de colonnes. Concaténez les jeux de données avec concatenate_datasets() :
>>> from datasets import concatenate_datasets, load_dataset
>>> stories = load_dataset("ajibawa-2023/General-Stories-Collection", split="train", streaming=True)
>>> stories = stories.select_columns(["text"]) # only keep the 'text' column
IterableDataset({
features: Unknown,
num_shards: 10
})
>>> wiki = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
>>> wiki = wiki.select_columns(["text"]) # only keep the 'text' column
IterableDataset({
features: ['text'],
num_shards: 41
})
>>> bert_dataset = concatenate_datasets([stories, wiki])
IterableDataset({
features: ['text'],
num_shards: 51
})Les shards du dataset concaténé sont les shards des datasets d'entrée.
Vous pouvez également concaténer deux jeux de données horizontalement en définissant axis=1, à condition que les jeux de données aient le même nombre de lignes :
>>> from datasets import IterableDataset
>>> stories_ids = IterableDataset.from_dict({"ids": list(range(num_stories))})
>>> stories_with_ids = concatenate_datasets([stories, stories_ids], axis=1)Dans ce cas, le jeu de données concaténé ne comporte qu'un seul shard, afin d'éviter de se retrouver avec des shards non alignés issus des jeux de données d'entrée.
Interleaver
interleave_datasets() peut combiner un IterableDataset avec d’autres jeux de données s’ils partagent les mêmes types de colonnes. Le jeu de données combiné renvoie des exemples alternés issus de chacun des jeux de données d’origine.
>>> from datasets import interleave_datasets
>>> es_dataset = load_dataset('allenai/c4', 'es', split='train', streaming=True)
IterableDataset({
features: Unknown,
num_shards: 2048
})
>>> fr_dataset = load_dataset('allenai/c4', 'fr', split='train', streaming=True)
IterableDataset({
features: Unknown,
num_shards: 2048
})
>>> multilingual_dataset = interleave_datasets([es_dataset, fr_dataset])
IterableDataset({
features: ['text', 'timestamp', 'url'],
num_shards: 2048
})
>>> list(multilingual_dataset.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
{'text': 'Le sacre de philippe ier, 23 mai 1059 - Compte Rendu...'}]Définissez les probabilités d'échantillonnage à partir de chacun des jeux de données d'origine pour un meilleur contrôle sur la manière dont chacun d'eux est échantillonné et combiné. Définissez l'argument probabilities avec les probabilités d'échantillonnage souhaitées :
>>> multilingual_dataset_with_oversampling = interleave_datasets([es_dataset, fr_dataset], probabilities=[0.8, 0.2], seed=42)
>>> list(multilingual_dataset_with_oversampling.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
{'text': 'Chevrolet Cavalier Usados en Bogota - Carros en Vent...'}]Environ 80 % du jeu de données final est composé du es_dataset, et 20 % du fr_dataset.
Vous pouvez également spécifier stopping_strategy. La stratégie par défaut, first_exhausted, est une stratégie de sous-échantillonnage, c'est-à-dire que la construction du jeu de données est arrêtée dès qu'un des jeux de données épuise ses échantillons.
Vous pouvez spécifier stopping_strategy=all_exhausted pour exécuter une stratégie de sur-échantillonnage. Dans ce cas, la construction du jeu de données est arrêtée dès que chaque échantillon de chaque jeu de données a été ajouté au moins une fois. En pratique, cela signifie qu'un jeu de données épuisé reviendra au début de ce jeu de données jusqu'à ce que le critère d'arrêt soit atteint.
Notez que si aucune probabilité d'échantillonnage n'est spécifiée, le nouveau jeu de données aura max_length_datasets*nb_dataset samples.
Il existe également stopping_strategy=all_exhausted_without_replacement pour s'assurer que chaque échantillon est vu exactement une fois.
Pour garantir un parallélisme correct à l'aide du sharding, les shards du jeu de données entrelacé contiennent au moins 1 shard de chaque jeu de données d'entrée. Par conséquent, le niveau de sharding du jeu de données entrelacé est le niveau de sharding minimal des jeux de données d'entrée.
Par exemple, si les jeux de données d’entrée ont respectivement 32, 48 et 128 shards, alors le jeu de données entrelacé a 32 = min(32, 48, 128) shards, et chaque nouveau shard contient 1 shard du premier jeu de données, 1 à 2 shards du deuxième jeu de données et 4 shards du troisième jeu de données.
Renommer, supprimer et convertir
Les méthodes suivantes vous permettent de modifier les colonnes d'un jeu de données. Ces méthodes sont utiles pour renommer ou supprimer des colonnes et pour convertir des colonnes en un nouvel ensemble de caractéristiques.
Renommer
Utilisez IterableDataset.rename_column() lorsque vous devez renommer une colonne de votre jeu de données. Les fonctionnalités associées à la colonne d’origine sont en réalité déplacées sous le nouveau nom de colonne, plutôt que de simplement remplacer la colonne d’origine sur place.
Fournissez à IterableDataset.rename_column() le nom de la colonne d’origine et le nom de la nouvelle colonne :
>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.rename_column("text", "content")Supprimer
Lorsque vous devez supprimer une ou plusieurs colonnes, transmettez à IterableDataset.remove_columns() le nom de la colonne à supprimer. Pour supprimer plusieurs colonnes, fournissez une liste de noms de colonnes :
>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.remove_columns('timestamp')Cast
IterableDataset.cast() modifie le type de fonction d’une ou de plusieurs colonnes. Cette méthode prend vos nouvelles Features en argument. L’exemple de code suivant montre comment modifier les types de fonction de ClassLabel et Value :
>>> from datasets import load_dataset
>>> dataset = load_dataset('nyu-mll/glue', 'mrpc', split='train', streaming=True)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['not_equivalent', 'equivalent']),
'idx': Value('int32')}
>>> from datasets import ClassLabel, Value
>>> new_features = dataset.features.copy()
>>> new_features["label"] = ClassLabel(names=['negative', 'positive'])
>>> new_features["idx"] = Value('int64')
>>> dataset = dataset.cast(new_features)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['negative', 'positive']),
'idx': Value('int64')}Le casting ne fonctionne que si le type de fonctionnalité d’origine et le nouveau type de fonctionnalité sont compatibles. Par exemple, vous pouvez caster une colonne avec le type de fonctionnalité
Value('int32')enValue('bool')si la colonne d’origine ne contient que des uns et des zéros.
Utilisez IterableDataset.cast_column() pour modifier le type de fonctionnalité d’une seule colonne. Passez le nom de la colonne et son nouveau type de fonctionnalité en arguments :
>>> dataset.features
{'audio': Audio(sampling_rate=44100, mono=True)}
>>> dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
>>> dataset.features
{'audio': Audio(sampling_rate=16000, mono=True)}Carte
De manière similaire à la fonction Dataset.map() pour un Dataset classique, 🤗 Datasets propose IterableDataset.map() pour traiter un IterableDataset. IterableDataset.map() applique le traitement à la volée lorsque les exemples sont diffusés.
Elle permet d'appliquer une fonction de traitement à chaque exemple d'un jeu de données, de manière indépendante ou par lots. Cette fonction peut même créer de nouvelles lignes et colonnes.
L'exemple suivant montre comment tokeniser un IterableDataset. La fonction doit accepter et renvoyer un dict :
>>> def add_prefix(example):
... example['text'] = 'My text: ' + example['text']
... return exampleEnsuite, appliquez cette fonction au jeu de données avec IterableDataset.map() :
>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> updated_dataset = dataset.map(add_prefix)
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...',
'timestamp': '2019-04-25 12:57:54',
'url': 'https://klyq.com/beginners-bbq-class-taking-place-in-missoula/'},
{'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...',
'timestamp': '2019-04-21 10:07:13',
'url': 'https://forums.macrumors.com/threads/restore-from-larger-disk-to-smaller-disk.1311329/'},
{'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...',
'timestamp': '2019-04-25 10:40:23',
'url': 'https://awishcometrue.com/Catalogs/Clearance/Tweens/V1960-Find-A-Way'}]Prenons un autre exemple, mais cette fois-ci, vous supprimerez des colonnes avec IterableDataset.map(). Lorsque vous supprimez une colonne, elle n’est retirée qu’après que l’exemple a été transmis à la fonction mappée. Cela permet à la fonction mappée d’utiliser le contenu des colonnes avant leur suppression.
Spécifiez la colonne à supprimer avec l'argument remove_columns dans IterableDataset.map() :
>>> updated_dataset = dataset.map(add_prefix, remove_columns=["timestamp", "url"])
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...'},
{'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...'},
{'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...'}]Traitement par lots
IterableDataset.map() prend également en charge le traitement par lots d’exemples. Pour opérer sur des lots, définissez batched=True. La taille de lot par défaut est de 1000, mais vous pouvez l’ajuster avec l’argument batch_size. Cela ouvre la porte à de nombreuses applications intéressantes, telles que la tokenisation, la découpe de longues phrases en segments plus courts et l’augmentation de données.
Tokenisation
>>> from datasets import load_dataset
>>> from transformers import AutoTokenizer
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
>>> def encode(examples):
... return tokenizer(examples['text'], truncation=True, padding='max_length')
>>> dataset = dataset.map(encode, batched=True, remove_columns=["text", "timestamp", "url"])
>>> next(iter(dataset))
{'input_ids': [101, 4088, 16912, 22861, 4160, 2465, 2635, 2173, 1999, 3335, ..., 0, 0, 0],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 0, 0]}Voir d'autres exemples de traitement par lots dans la documentation batched map processing. Ils fonctionnent de la même manière pour les jeux de données itérables.
Filtre
Vous pouvez filtrer les lignes du jeu de données en fonction d'une fonction prédicat à l'aide de Dataset.filter(). Elle renvoie les lignes qui correspondent à une condition spécifiée :
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', streaming=True, split='train')
>>> start_with_ar = dataset.filter(lambda example: example['text'].startswith('San Francisco'))
>>> next(iter(start_with_ar))
{'text': 'San Francisco 49ers cornerback Shawntae Spencer will miss the rest of the sea...}Dataset.filter() peut également filtrer par indices si vous définissez with_indices=True :
>>> even_dataset = dataset.filter(lambda example, idx: idx % 2 == 0, with_indices=True)
>>> list(even_dataset.take(3))
[{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012 Whe...},
{'text': 'Car Wash For Clara! Now is your chance to help! 2 year old Clara Woodward has...},
{'text': 'Log In Please enter your ECode to log in. Forgotten your eCode? If you create...}]Lot
La méthode batch transforme votre IterableDataset en un itérable de lots. C'est particulièrement utile lorsque vous souhaitez travailler avec des lots dans votre boucle d'entraînement ou lorsque vous utilisez des frameworks qui s'attendent à des entrées par lots.
Il existe également une option « Batch Processing » lors de l'utilisation de la fonction
mappour appliquer une fonction à des lots de données, ce qui est abordé dans la section Map ci-dessus. La méthodebatchdécrite ici est différente et offre un moyen plus direct de créer des lots à partir de votre jeu de données.
Vous pouvez utiliser la méthode batch de cette manière :
from datasets import load_dataset
# Load a dataset in streaming mode
dataset = load_dataset("some_dataset", split="train", streaming=True)
# Create batches of 32 samples
batched_dataset = dataset.batch(batch_size=32)
# Iterate over the batched dataset
for batch in batched_dataset:
print(batch)
breakDans cet exemple, batched_dataset est toujours un IterableDataset, mais chaque élément produit est désormais un lot de 32 échantillons au lieu d'un seul échantillon. Ce groupement par lots est effectué à la volée lors de l'itération sur le jeu de données, préservant ainsi la nature économe en mémoire de IterableDataset.
La méthode par lots fournit également un paramètre drop_last_batch. Lorsqu’il est défini sur True, elle supprime le dernier lot s’il est plus petit que la batch_size spécifiée. Cela peut être utile dans les scénarios où votre traitement en aval exige que tous les lots aient la même taille :
batched_dataset = dataset.batch(batch_size=32, drop_last_batch=True)Stream dans une boucle d'entraînement
IterableDataset peut être intégré dans une boucle d’entraînement. Commencez par mélanger le jeu de données :
>>> seed, buffer_size = 42, 10_000
>>> dataset = dataset.shuffle(seed, buffer_size=buffer_size)Enfin, créez une boucle d’entraînement simple et lancez l’entraînement :
>>> import torch
>>> from torch.utils.data import DataLoader
>>> from transformers import AutoModelForMaskedLM, DataCollatorForLanguageModeling
>>> from tqdm import tqdm
>>> dataset = dataset.with_format("torch")
>>> dataloader = DataLoader(dataset, collate_fn=DataCollatorForLanguageModeling(tokenizer))
>>> device = 'cuda' if torch.cuda.is_available() else 'cpu'
>>> model = AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased")
>>> model.train().to(device)
>>> optimizer = torch.optim.AdamW(params=model.parameters(), lr=1e-5)
>>> for epoch in range(3):
... dataset.set_epoch(epoch)
... for i, batch in enumerate(tqdm(dataloader, total=5)):
... if i == 5:
... break
... batch = {k: v.to(device) for k, v in batch.items()}
... outputs = model(**batch)
... loss = outputs[0]
... loss.backward()
... optimizer.step()
... optimizer.zero_grad()
... if i % 10 == 0:
... print(f"loss: {loss}")Enregistrer un point de contrôle de jeu de données et reprendre l'itération
Si votre boucle d’entraînement s’arrête, vous souhaiterez peut-être reprendre l’entraînement là où il s’était interrompu. Pour ce faire, vous pouvez enregistrer un point de contrôle de votre modèle et de vos optimiseurs, ainsi que de votre chargeur de données.
Les jeux de données itérables ne permettent pas d’accéder de manière aléatoire à un index d’exemple spécifique pour reprendre la lecture, mais vous pouvez utiliser IterableDataset.state_dict() et IterableDataset.load_state_dict() pour reprendre depuis un point de contrôle, de la même manière que vous le feriez pour les modèles et les optimiseurs :
>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> for idx, example in enumerate(iterable_dataset):
... print(example)
... if idx == 2:
... state_dict = iterable_dataset.state_dict()
... print("checkpoint")
... break
>>> iterable_dataset.load_state_dict(state_dict)
>>> print(f"restart from checkpoint")
>>> for example in iterable_dataset:
... print(example)Retourne :
{'a': 0}
{'a': 1}
{'a': 2}
checkpoint
restart from checkpoint
{'a': 3}
{'a': 4}
{'a': 5}En interne, le jeu de données itérable suit le shard en cours de lecture et l’index de l’exemple dans le shard actuel, et il stocke ces informations dans le state_dict.
Pour reprendre depuis un point de contrôle, le jeu de données ignore tous les fragments précédemment lus afin de redémarrer depuis le fragment actuel. Il lit ensuite le fragment et ignore les exemples jusqu'à atteindre l'exemple exact du point de contrôle.
La relance d’un jeu de données est donc assez rapide, car il ne relit pas les shards déjà parcourus. Toutefois, la reprise d’un jeu de données n’est généralement pas instantanée, car il faut relire depuis le début du shard en cours et ignorer les exemples jusqu’à atteindre l’emplacement du point de contrôle.
Cela peut être utilisé avec le StatefulDataLoader de torchdata :
>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(iterable_dataset, batch_size=32, num_workers=4)
>>> # checkpoint
>>> state_dict = dataloader.state_dict() # uses iterable_dataset.state_dict() under the hood
>>> # resume from checkpoint
>>> dataloader.load_state_dict(state_dict) # uses iterable_dataset.load_state_dict() under the hoodLa reprise reprend exactement là où le point de contrôle a été enregistré, sauf si
.shuffle()est utilisé : les exemples des tampons de mélange sont perdus lors de la reprise et les tampons sont remplis à nouveau avec de nouvelles données.
Enregistrer
Une fois votre jeu de données itérable prêt, vous pouvez l'enregistrer au format Parquet en tant que jeu de données Hugging Face et le réutiliser plus tard avec load_dataset().
Enregistrez votre jeu de données en fournissant le nom du dépôt de jeu de données sur Hugging Face où vous souhaitez l’enregistrer en tant qu’argument de push_to_hub(). Cela parcourt le jeu de données et télécharge progressivement les données sur Hugging Face :
dataset.push_to_hub("username/my_dataset")Par défaut, IterableDataset.push_to_hub() écrit dataset.num_shards fichiers. Définissez num_shards pour choisir un nombre différent. S'il y a plus de shards de sortie que de dataset.num_shards, chaque shard d'entrée est réparti sur son propre groupe de shards de sortie en y écrivant ses exemples de manière circulaire (round-robin), en une seule passe sur les données. S'il y a moins d'exemples que de shards demandés, les shards restants sont des fichiers Parquet vides.
Si le jeu de données est composé de plusieurs fragments (dataset.num_shards > 1), vous pouvez utiliser plusieurs processus pour le téléverser en parallèle. Cela est particulièrement utile si vous avez appliqué des étapes map() ou filter(), car elles s'exécuteront plus rapidement en parallèle :
dataset.push_to_hub("username/my_dataset", num_proc=8)Utilisez la fonction load_dataset() pour recharger le jeu de données :
from datasets import load_dataset
reloaded_dataset = load_dataset("username/my_dataset")Export
🤗 Datasets prend en charge l'exportation, vous permettant ainsi de travailler avec votre jeu de données dans d'autres applications. Le tableau suivant présente les formats de fichiers actuellement pris en charge pour l'exportation :
| Type de fichier | Méthode d'export |
|---|---|
| CSV | IterableDataset.to_csv() |
| JSON | IterableDataset.to_json() |
| Parquet | IterableDataset.to_parquet() |
| SQL | IterableDataset.to_sql() |
| Objet Python en mémoire | IterableDataset.to_pandas(), IterableDataset.to_polars() or IterableDataset.to_dict() |
Par exemple, exportez votre jeu de données dans un fichier CSV comme suit :
>>> dataset.to_csv("path/of/my/dataset.csv")Si vous disposez d'un grand jeu de données, vous pouvez enregistrer un fichier par shard, par exemple
>>> num_shards = dataset.num_shards
>>> for index in range(num_shards):
... shard = dataset.shard(index, num_shards)
... shard.to_parquet(f"path/of/my/dataset/data-{index:05d}.parquet")