Dask

Dask é uma biblioteca de computação paralela e distribuída que escala o ecossistema existente de Python e PyData.

Em particular, podemos usar Dask DataFrame para escalar fluxos de trabalho do pandas. O Dask DataFrame paraleliza o pandas para lidar com grandes dados tabulares. Ele espelha de perto a API do pandas, tornando simples a transição de testes em um único conjunto de dados para o processamento do conjunto completo. O Dask é particularmente eficaz com Parquet, o formato padrão no Hugging Face Datasets, pois suporta tipos de dados ricos, filtragem colunar eficiente e compressão.

Um bom caso de uso prático para o Dask é executar o processamento de dados ou a inferência de modelos em um conjunto de dados de forma distribuída. Veja, por exemplo, o excelente post de blog da Coiled sobre Scaling AI-Based Data Processing with Hugging Face + Dask.

Ler e Escrever

Como o Dask usa fsspec para ler e gravar dados remotos, você pode usar os caminhos do Hugging Face (hf://) para ler e gravar dados no Hub.

Primeiro, você precisa fazer login com sua conta do Hugging Face, por exemplo, usando:

hf auth login

Em seguida, você pode criar um repositório de dataset, por exemplo, usando:

from huggingface_hub import HfApi

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

Por fim, você pode usar caminhos do Hugging Face no Dask. O Dask DataFrame suporta escrita distribuída em Parquet no Hugging Face, que usa commits para rastrear alterações no dataset:

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

Como isso cria um commit por arquivo, é recomendado compactar o histórico após o upload:

from huggingface_hub import HfApi

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

Isso cria um repositório de dataset username/my_dataset contendo seu dataset Dask em formato Parquet. Você pode carregá-lo novamente mais tarde:

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 mais informações sobre os caminhos do Hugging Face e como eles são implementados, consulte a documentação da biblioteca de cliente sobre o HfFileSystem.

Processar dados

Para processar um conjunto de dados em paralelo usando Dask, você pode primeiro definir sua função de processamento de dados para um DataFrame ou Series do pandas e, em seguida, usar a função map_partitions do Dask para aplicar essa função a todas as partições de um conjunto de dados em paralelo:

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

ou uma função semelhante usando os métodos de string do pandas (mais rápido):

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

No pandas, você pode usar essa função em uma coluna de texto:

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

E no Dask você pode executar essa função em cada partição:

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

Observe que você também precisa fornecer meta, que é o tipo da Series ou DataFrame do pandas na saída da sua função. Isso é necessário porque o Dask DataFrame usa uma API preguiçosa. Como o Dask só executará o processamento de dados quando .compute() for chamado, ele precisa do argumento meta para conhecer o tipo da nova coluna no meio tempo.

Pushdown de predicados e projeções

Ao ler dados Parquet do Hugging Face, o Dask aproveita automaticamente os metadados nos arquivos Parquet para ignorar arquivos ou grupos de linhas inteiros, caso não sejam necessários. Por exemplo, se você aplicar um filtro (predicado) em um Dataset do Hugging Face em formato Parquet ou selecionar um subconjunto de colunas (projeção), o Dask lerá os metadados dos arquivos Parquet para descartar as partes que não são necessárias, sem baixá-las.

Isso é possível graças a uma reimplementação da API do Dask DataFrame para suportar otimização de consultas, o que torna o Dask mais rápido e robusto.

Por exemplo, esse subconjunto do FineWeb-Edu contém muitos arquivos Parquet. Se você puder filtrar o dataset para manter o texto de dumps recentes do CC, o Dask ignorará a maioria dos arquivos e baixará apenas os dados que correspondem ao 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"]

O Dask também lerá apenas as colunas necessárias para o seu cálculo e ignorará as demais. Por exemplo, se você remover uma coluna mais tarde no seu código, ele não se dará ao trabalho de carregá-la no início do pipeline, caso não seja necessária. Isso é útil quando você deseja manipular um subconjunto das colunas ou para análises:

# 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

A maioria dos recursos em dask é otimizada para um cluster ou um Client local para iniciar os cálculos paralelos:

import dask.dataframe as dd
from distributed import Client

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

Para uso local, o Client usa um LocalCluster do Dask com multiprocessamento por padrão. Você pode configurar manualmente o multiprocessamento do LocalCluster com

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

Observe que, se você usar o agendador de threads padrão localmente sem Client, um DataFrame pode ficar mais lento após determinadas operações (mais detalhes aqui).

Encontre mais informações sobre a configuração de um cluster local ou na nuvem na documentação de implantação do Dask.