Dask

Dask est une bibliothèque de calcul parallèle et distribué qui étend l'écosystème Python et PyData existant.

En particulier, nous pouvons utiliser Dask DataFrame pour mettre à l'échelle les flux de travail pandas. Dask DataFrame parallélise pandas pour traiter de grands jeux de données tabulaires. Il reproduit fidèlement l'API de pandas, ce qui facilite la transition des tests sur un seul jeu de données au traitement du jeu de données complet. Dask est particulièrement efficace avec Parquet, le format par défaut sur Hugging Face Datasets, car il prend en charge des types de données riches, un filtrage columnaire efficace et la compression.

Un bon cas d’usage pratique de Dask consiste à exécuter le traitement de données ou l’inférence de modèles sur un jeu de données de manière distribuée. Consultez, par exemple, l’excellent article de blog de Coiled sur Scaling AI-Based Data Processing with Hugging Face + Dask.

Lecture et écriture

Puisque Dask utilise fsspec pour lire et écrire des données distantes, vous pouvez utiliser les chemins Hugging Face (hf://) pour lire et écrire des données sur le Hub.

Tout d'abord, vous devez vous connecter avec votre compte Hugging Face, par exemple en utilisant :

hf auth login

Ensuite, vous pouvez Créer un dépôt de jeu de données, par exemple en utilisant :

from huggingface_hub import HfApi

HfApi().create_repo(repo_id="username/my_dataset", repo_type="dataset")

Enfin, vous pouvez utiliser des chemins Hugging Face dans Dask. Dask DataFrame prend en charge l'écriture distribuée vers Parquet sur Hugging Face, ce qui utilise des commits pour suivre les modifications du jeu de données :

import dask.dataframe as dd

df.to_parquet("hf://datasets/username/my_dataset")

# or write in separate directories if the dataset has train/validation/test splits
df_train.to_parquet("hf://datasets/username/my_dataset/train")
df_valid.to_parquet("hf://datasets/username/my_dataset/validation")
df_test .to_parquet("hf://datasets/username/my_dataset/test")

Comme cela crée un commit par fichier, il est recommandé d'écraser l'historique après le téléversement :

from huggingface_hub import HfApi

HfApi().super_squash_history(repo_id=repo_id, repo_type="dataset")

Cela crée un dépôt de jeu de données username/my_dataset contenant votre jeu de données Dask au format Parquet. Vous pouvez le recharger plus tard :

import dask.dataframe as dd

df = dd.read_parquet("hf://datasets/username/my_dataset")

# or read from separate directories if the dataset has train/validation/test splits
df_train = dd.read_parquet("hf://datasets/username/my_dataset/train")
df_valid = dd.read_parquet("hf://datasets/username/my_dataset/validation")
df_test  = dd.read_parquet("hf://datasets/username/my_dataset/test")

Pour plus d’informations sur les chemins Hugging Face et leur implémentation, veuillez consulter la documentation de la bibliothèque client sur HfFileSystem.

Traiter les données

Pour traiter un jeu de données en parallèle à l'aide de Dask, vous pouvez d'abord définir votre fonction de traitement de données pour un DataFrame ou une Series pandas, puis utiliser la fonction map_partitions de Dask pour appliquer cette fonction à toutes les partitions d'un jeu de données en parallèle :

def dummy_count_words(texts):
    return pd.Series([len(text.split(" ")) for text in texts])

ou une fonction similaire utilisant les méthodes de chaîne de pandas (plus rapide) :

def dummy_count_words(texts):
    return texts.str.count(" ")

Dans pandas, vous pouvez utiliser cette fonction sur une colonne de texte :

# pandas API
df["num_words"] = dummy_count_words(df.text)

Et dans Dask, vous pouvez exécuter cette fonction sur chaque partition :

# Dask API: run the function on every partition
df["num_words"] = df.text.map_partitions(dummy_count_words, meta=int)

Notez que vous devez également fournir meta, qui est le type de la Series ou du DataFrame pandas dans la sortie de votre fonction. Ceci est nécessaire car Dask DataFrame utilise une API paresseuse. Comme Dask n’exécutera le traitement des données qu’une fois .compute() appelé, il a besoin de l’argument meta pour connaître le type de la nouvelle colonne en attendant.

Poussée des prédicats et des projections

Lors de la lecture de données Parquet depuis Hugging Face, Dask exploite automatiquement les métadonnées des fichiers Parquet pour ignorer des fichiers ou des groupes de lignes entiers s’ils ne sont pas nécessaires. Par exemple, si vous appliquez un filtre (prédicat) à un Dataset Hugging Face au format Parquet ou si vous sélectionnez un sous-ensemble de colonnes (projection), Dask lit les métadonnées des fichiers Parquet pour écarter les parties non nécessaires sans les télécharger.

Cela est rendu possible grâce à une réimplémentation de l'API Dask DataFrame pour prendre en charge l'optimisation des requêtes, ce qui rend Dask plus rapide et plus robuste.

Par exemple, ce sous-ensemble de FineWeb-Edu contient de nombreux fichiers Parquet. Si vous pouvez filtrer le jeu de données pour conserver le texte des dumps CC récents, Dask ignorera la plupart des fichiers et ne téléchargera que les données correspondant au filtre :

import dask.dataframe as dd

df = dd.read_parquet("hf://datasets/HuggingFaceFW/fineweb-edu/sample/10BT/*.parquet")

# Dask will skip the files or row groups that don't
# match the query without downloading them.
df = df[df.dump >= "CC-MAIN-2023"]

Dask ne lira également que les colonnes requises pour votre calcul et ignorera les autres. Par exemple, si vous supprimez une colonne tardivement dans votre code, il ne se chargera pas de la charger en début de pipeline si elle n’est pas nécessaire. Cela est utile lorsque vous souhaitez manipuler un sous-ensemble de colonnes ou pour l’analyse de données :

# Dask will download the 'dump' and 'token_count' needed
# for the filtering and computation and skip the other columns.
df.token_count.mean().compute()

Client

La plupart des fonctionnalités de dask sont optimisées pour un cluster ou un Client local afin de lancer les calculs parallèles :

import dask.dataframe as dd
from distributed import Client

if __name__ == "__main__":  # needed for creating new processes
    client = Client()
    df = dd.read_parquet(...)
    ...

Pour un usage local, le Client utilise par défaut un LocalCluster Dask avec le multiprocessing. Vous pouvez configurer manuellement le multiprocessing de LocalCluster avec

from dask.distributed import Client, LocalCluster
cluster = LocalCluster(n_workers=8, threads_per_worker=8)
client = Client(cluster)

Notez que si vous utilisez le planificateur par défaut à base de threads localement sans Client, un DataFrame peut devenir plus lent après certaines opérations (plus de détails ici).

Retrouvez plus d’informations sur la configuration d’un cluster local ou cloud dans la documentation du déploiement de Dask.