Bibliothèque Datasets de Hugging Face : Charger et Traiter

1. Aperçu général

Quand vous travaillez avec un modèle NLP ou un fine-tuning LLM, vous avez toujours besoin d’une étape préalable : préparer les données au bon format, suffisamment propres, et pouvoir les charger dans le pipeline d’entraînement. La bibliothèque datasets de Hugging Face résout exactement ce problème — elle vous fournit une API unifiée pour charger des datasets depuis le Hub, traiter par lots avec map/filter, et les sauvegarder dans un format optimisé pour une lecture répétée.

Après avoir lu cet article, vous serez capable de charger des datasets depuis le Hub Hugging Face, de comprendre la structure Dataset/DatasetDict, de savoir comment prétraiter le texte avec un tokenizer, de traiter les grands datasets en streaming, et d’écrire des tests pour votre pipeline de traitement des données.

2. Installation

pip install datasets transformers

La bibliothèque datasets ne dépend pas strictement de transformers, mais la partie tokenization au point 6 en a besoin.

import datasets
print(datasets.__version__)
3.0.1

3. Charger des datasets depuis le Hub Hugging Face

3.1. load_dataset basique

from datasets import load_dataset

reviews = load_dataset("imdb")
print(reviews)
DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    unsupervised: Dataset({
        features: ['text', 'label'],
        num_rows: 50000
    })
})

load_dataset charge le script/métadonnées du dataset imdb depuis le Hub, met en cache le fichier compressé localement (~/.cache/huggingface/datasets), puis retourne un DatasetDict — c’est-à-dire un dictionnaire de Dataset selon les splits (train, test, unsupervised).

3.2. Ne récupérer qu’un split spécifique

train_reviews = load_dataset("imdb", split="train")
print(type(train_reviews), len(train_reviews))
<class 'datasets.arrow_dataset.Dataset'> 25000

Passer split="train" retourne directement un Dataset au lieu d’un DatasetDict, ce qui est pratique quand vous n’avez besoin que d’un seul ensemble.

Remarque: Si le dataset sur le Hub nécessite une config (par exemple glue), vous devez passer le nom de la config en plus : load_dataset("glue", "mrpc"). Omettre ce paramètre lèvera une erreur demandant de choisir une config.

4. Explorer la structure d’un Dataset

sample = train_reviews[0]
print(sample.keys())
print(sample["label"], sample["text"][:60])
dict_keys(['text', 'label'])
0 I rented I AM CURIOUS-YELLOW from my video store because of...

L’accès par index retourne un dictionnaire Python ordinaire. Dataset supporte aussi le slicing qui retourne un dictionnaire de listes :

batch = train_reviews[0:3]
print(batch["label"])
[0, 0, 0]

Voir le schéma via features :

print(train_reviews.features)
{'text': Value(dtype='string', id=None),
 'label': ClassLabel(names=['neg', 'pos'], id=None)}

ClassLabel vous indique que 0 correspond au label neg, 1 est pos — utile pour une correspondance inverse vers les noms de labels.

5. Traiter les données avec map, filter, select

5.1. filter pour filtrer les lignes

positive_only = train_reviews.filter(lambda row: row["label"] == 1)
print(len(positive_only))
12500

filter accepte une fonction qui retourne True/False pour chaque ligne, parcourt tout le dataset et ne conserve que les lignes qui satisfont la condition.

5.2. map pour transformer les données

def add_length(row):
    row["text_length"] = len(row["text"])
    return row

train_with_length = train_reviews.map(add_length)
print(train_with_length[0]["text_length"])
1640

map applique la fonction à chaque ligne, retourne un dictionnaire contenant les nouveaux champs (ou mise à jour des anciens) — le résultat est fusionné dans le dataset d’origine.

5.3. map par batch pour accélérer

def add_length_batched(batch):
    batch["text_length"] = [len(t) for t in batch["text"]]
    return batch

train_fast = train_reviews.map(add_length_batched, batched=True, batch_size=1000)

Avec batched=True, la fonction reçoit un dictionnaire batch entier (chaque champ est une liste) au lieu de chaque ligne individuellement — cela évite d’appeler la fonction Python 25 000 fois, beaucoup plus rapide avec les grands datasets.

5.4. select pour extraire un sous-ensemble

small_train = train_reviews.select(range(100))
print(len(small_train))
100

À utiliser quand vous voulez tester rapidement votre pipeline sur un sous-ensemble avant d’exécuter le dataset complet.

6. Prétraitement du texte avec un tokenizer

Le cas le plus courant en pratique : tokenizer le texte pour le passer dans le modèle.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_batch(batch):
    return tokenizer(batch["text"], truncation=True, max_length=256)

tokenized = train_reviews.map(tokenize_batch, batched=True)
print(tokenized.column_names)
['text', 'label', 'input_ids', 'attention_mask']

tokenize_batch reçoit un batch et retourne un dictionnaire contenant input_ids, attention_mask — map fusionne automatiquement ces nouveaux champs dans le dataset, en conservant les anciens champs. Pour supprimer les champs inutiles (par exemple text brut pour économiser la RAM), utilisez remove_columns :

tokenized = train_reviews.map(
    tokenize_batch, batched=True, remove_columns=["text"]
)
print(tokenized.column_names)
['label', 'input_ids', 'attention_mask']

7. Diviser, mélanger, fragmenter les splits

7.1. Diviser soi-même train/test

Si le dataset n’a qu’un seul split, vous pouvez vous-même le diviser avec train_test_split :

split_result = train_reviews.train_test_split(test_size=0.2, seed=42)
print(split_result)
DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 20000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 5000
    })
})

seed=42 garantit que la division est reproductible entre les exécutions.

7.2. Mélanger et fragmenter

shuffled = train_reviews.shuffle(seed=42)
part = shuffled.shard(num_shards=4, index=0)
print(len(part))
6250

shard divise le dataset en num_shards parts égales — utile quand vous distribuez les données à plusieurs workers d’entraînement distribué.

8. Streaming pour les grands datasets

Quand le dataset est trop grand pour être chargé entièrement en mémoire/disque, utilisez streaming=True :

streamed = load_dataset("imdb", split="train", streaming=True)
first_three = list(streamed.take(3))
print(len(first_three))
3

Avec le streaming, datasets ne télécharge pas le fichier entier mais lit séquentiellement chaque enregistrement via un itérateur. Vous ne pouvez pas utiliser len() ou l’indexage direct, mais vous pouvez toujours map/filter de manière paresseuse (appliqué seulement lors de l’itération).

Remarque: Le streaming convient aux datasets de centaines de GB comme les ensembles de préentraînement. Avec les datasets petits (quelques centaines de MB), le chargement normal (streaming=False) est plus rapide car il exploite le cache et l’accès aléatoire.

9. Sauvegarder et recharger un dataset

Après avoir terminé le traitement, sauvegardez le résultat pour ne pas avoir à relancer map à chaque fois :

tokenized.save_to_disk("data/imdb_tokenized")
from datasets import load_from_disk

reloaded = load_from_disk("data/imdb_tokenized")
print(len(reloaded))
25000

save_to_disk écrit le dataset au format Arrow — relire est très rapide car il n’y a pas besoin de reanalyser à partir de la source. Si vous avez besoin d’interagir avec pandas, convertissez directement :

df = tokenized.select(range(5)).to_pandas()
print(df.columns.tolist())
['label', 'input_ids', 'attention_mask']

10. Écrire des tests pour votre pipeline de traitement des données

Le pipeline de traitement des données devrait avoir des tests pour garantir que la logique map/filter ne se casse pas silencieusement quand vous modifiez le code. Voici un exemple de test avec pytest.

import pytest
from datasets import Dataset


def add_length(row):
    row["text_length"] = len(row["text"])
    return row


def test_add_length_computes_correct_value():
    data = Dataset.from_dict({"text": ["hello", "hi there"]})

    result = data.map(add_length)

    assert result[0]["text_length"] == 5
    assert result[1]["text_length"] == 8


def test_filter_keeps_only_positive_labels():
    data = Dataset.from_dict({"label": [0, 1, 1, 0]})

    positive = data.filter(lambda row: row["label"] == 1)

    assert len(positive) == 2
    assert all(row["label"] == 1 for row in positive)

Dataset.from_dict crée un petit dataset en mémoire, pas besoin de charger depuis le Hub — parfait pour les tests unitaires. Les deux tests ci-dessus vérifient le bon fonctionnement de map (génère un nouveau champ) et filter (filtre selon la condition) indépendamment du dataset réel sur Internet.

pytest test_dataset_pipeline.py -v
test_dataset_pipeline.py::test_add_length_computes_correct_value PASSED
test_dataset_pipeline.py::test_filter_keeps_only_positive_labels PASSED

11. Erreurs courantes et quand ne pas utiliser

  • Oublier batched=True lors du traitement de grands datasets : map sans batch fonctionne lentement car il appelle la fonction Python pour chaque ligne une par une. Avec des centaines de milliers de lignes, préférez toujours batched=True si la logique le permet.
  • Utiliser un indexage aléatoire sur un dataset en streaming : streamed[5] lèvera une erreur car le dataset en streaming est un itérateur qui ne supporte pas l’accès aléatoire. Utilisez seulement .take(), .skip(), ou itérez séquentiellement.
  • Ne pas définir seed avec shuffle/train_test_split : les résultats ne seront pas reproductibles entre les exécutions, ce qui rend le débogage difficile lors de la comparaison des métriques entre les expériences.
  • Cache qui explose sans contrôle : chaque fois que vous faites un map avec une fonction différente, datasets écrit un nouveau fichier cache. Si votre machine manque d’espace disque, utilisez load_dataset(..., keep_in_memory=True) ou supprimez périodiquement ~/.cache/huggingface/datasets.
  • Ne pas utiliser datasets pour des données relationnelles complexes (plusieurs tables à joindre, nécessitant des transactions). C’est une bibliothèque optimisée pour les données tabulaires/texte plates utilisées en ML, pas un remplacement de base de données.

12. Conclusion

La bibliothèque datasets vous fournit une API élégante pour aller de « charger un dataset brut depuis le Hub » à « dataset tokenisé, prêt à entraîner » en seulement quelques commandes load_dataset, map, filter. Avec les petits datasets, vous chargez normalement et exploitez le cache Arrow ; avec les datasets énormes, passez à streaming=True pour ne pas avoir à tout télécharger sur le disque. Quand vous construisez un pipeline de traitement des données pour un vrai projet, n’oubliez pas d’écrire des tests pour vos fonctions map/filter en utilisant Dataset.from_dict — cela vous aide à détecter rapidement si la logique de prétraitement change accidentellement.

huggingfacedatasetspythonnlpmachine-learning