Dask

Dask es una biblioteca de computación paralela y distribuida que escala el ecosistema existente de Python y PyData.

En particular, podemos usar Dask DataFrame para escalar flujos de trabajo de pandas. Dask DataFrame paraleliza pandas para manejar grandes volúmenes de datos tabulares. Su API refleja de cerca la de pandas, lo que facilita la transición desde la prueba con un único conjunto de datos hasta el procesamiento del conjunto completo. Dask es especialmente eficaz con Parquet, el formato predeterminado en Hugging Face Datasets, ya que admite tipos de datos ricos, filtrado columnar eficiente y compresión.

Un buen caso de uso práctico de Dask es ejecutar el procesamiento de datos o la inferencia de modelos sobre un conjunto de datos de forma distribuida. Consulte, por ejemplo, la excelente publicación de Coiled sobre Scaling AI-Based Data Processing with Hugging Face + Dask.

Lectura y escritura

Dado que Dask utiliza fsspec para leer y escribir datos remotos, puedes usar las rutas de Hugging Face (hf://) para leer y escribir datos en el Hub.

Primero necesitas iniciar sesión con tu cuenta de Hugging Face, por ejemplo, usando:

hf auth login

Luego puedes crear un repositorio de datos, por ejemplo, usando:

from huggingface_hub import HfApi

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

Por último, puedes usar rutas de Hugging Face en Dask. Dask DataFrame admite la escritura distribuida en Parquet en Hugging Face, que utiliza commits para rastrear los cambios en el conjunto de datos:

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")

Dado que esto crea un commit por archivo, se recomienda comprimir el historial después de la carga:

from huggingface_hub import HfApi

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

Esto crea un repositorio de datos username/my_dataset que contiene tu conjunto de datos de Dask en formato Parquet. Puedes recargarlo más adelante:

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")

Para obtener más información sobre las rutas de Hugging Face y cómo se implementan, consulte la documentación de la biblioteca de cliente sobre HfFileSystem.

Procesar datos

Para procesar un conjunto de datos en paralelo con Dask, primero puede definir su función de procesamiento de datos para un DataFrame o Series de pandas y luego usar la función map_partitions de Dask para aplicar esta función a todas las particiones de un conjunto de datos en paralelo:

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

o una función similar que utilice los métodos de cadena de pandas (más rápido):

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

En pandas puedes usar esta función en una columna de texto:

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

Y en Dask puedes ejecutar esta función en cada partición:

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

Tenga en cuenta que también debe proporcionar meta, que es el tipo de la serie o el DataFrame de pandas en la salida de su función. Esto es necesario porque Dask DataFrame utiliza una API perezosa. Dado que Dask solo ejecutará el procesamiento de datos una vez que se llame a .compute(), necesita el argumento meta para conocer el tipo de la nueva columna mientras tanto.

Predicado y proyección

Al leer datos Parquet desde Hugging Face, Dask aprovecha automáticamente los metadatos de los archivos Parquet para omitir archivos completos o grupos de filas si no son necesarios. Por ejemplo, si se aplica un filtro (predicado) a un Dataset de Hugging Face en formato Parquet o si se selecciona un subconjunto de columnas (proyección), Dask leerá los metadatos de los archivos Parquet para descartar las partes que no se necesitan sin descargarlas.

Esto es posible gracias a una reimplementación de la API de DataFrames de Dask para admitir la optimización de consultas, lo que hace que Dask sea más rápido y robusto.

Por ejemplo, este subconjunto de FineWeb-Edu contiene muchos archivos Parquet. Si puedes filtrar el conjunto de datos para conservar el texto de volcados recientes de CC, Dask omitirá la mayoría de los archivos y solo descargará los datos que coincidan con el filtro:

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 también solo leerá las columnas necesarias para su cálculo y omitirá el resto. Por ejemplo, si elimina una columna tarde en su código, no se molestará en cargarla al principio de la tubería si no es necesaria. Esto es útil cuando desea manipular un subconjunto de columnas o para análisis:

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

Cliente

La mayoría de las funciones de dask están optimizadas para un clúster o un Client local para lanzar los cálculos en paralelo:

import dask.dataframe as dd
from distributed import Client

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

Para el uso local, el Client utiliza un LocalCluster de Dask con multiprocesamiento de forma predeterminada. Puedes configurar manualmente el multiprocesamiento de LocalCluster con

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

Tenga en cuenta que si utiliza el programador de hilos predeterminado localmente sin Client, un DataFrame puede volverse más lento después de ciertas operaciones (más detalles aquí).

Encuentra más información sobre la configuración de un clúster local o en la nube en la documentación de despliegue de Dask.