Dask

Daskは、既存のPythonおよびPyDataエコシステムをスケーリングするための並列・分散コンピューティングライブラリです。

特に、Dask DataFrame を使用して pandas ワークフローをスケールアップできます。Dask DataFrame は pandas を並列化し、大規模な表形式データを処理します。pandas の API に密接に準拠しているため、単一のデータセットでのテストからフルデータセットの処理への移行が容易です。Dask は、Hugging Face Datasets の既定の形式である Parquet と特に相性が良く、豊富なデータ型、効率的なカラム単位フィルタリング、および圧縮をサポートしています。

Daskの実用的なユースケースとして、データセットに対するデータ処理やモデル推論を分散方式で実行することが挙げられます。例えば、CoiledによるHugging Face + DaskによるAIベースのデータ処理のスケーリングに関する優れたブログ記事をご覧ください。

読み書き

Dask はリモートデータの読み書きに fsspec を使用するため、Hugging Face パス (hf://) を使って Hub 上のデータの読み書きを行うことができます。

まず、Hugging Face アカウントでログインする必要があります。例えば、以下のようにします:

hf auth login

その後、データセットリポジトリを作成できます。例えば、以下のように使用します:

from huggingface_hub import HfApi

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

最後に、Dask でHugging Face パスを使用できます。 Dask DataFrame は Hugging Face への Parquet の分散書き込みをサポートしており、コミットを使用してデータセットの変更を追跡します:

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

この操作はファイルごとに1つのコミットを作成するため、アップロード後に履歴をsquashすることをお勧めします。

from huggingface_hub import HfApi

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

これにより、Parquet形式のDaskデータセットを含むデータセットリポジトリusername/my_datasetが作成されます。 後で再読み込みできます:

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

Hugging Face のパスとその実装方法に関する詳細については、HfFileSystem に関するクライアントライブラリのドキュメントを参照してください。

データを処理する

Dask を使用してデータセットを並列処理するには、まず pandas の DataFrame または Series に対するデータ処理関数を定義し、その後 Dask の map_partitions 関数を使用して、この関数をデータセットのすべてのパーティションに並列で適用します。

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

または、pandas の文字列メソッドを使用した同様の関数（より高速）:

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

pandasでは、この関数をテキスト列に使用できます:

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

また、Daskでは、この関数をすべてのパーティションで実行できます:

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

出力される pandas の Series または DataFrame の型であるmetaも指定する必要があります。 これは、Dask DataFrame が遅延評価 API を使用しているためです。Dask は.compute()が呼び出されるまでデータ処理を実行しないため、その間、新しい列の型を知るためにmeta引数が必要です。

述語と射影のプッシュダウン

Hugging Face から Parquet データを読み込む場合、Dask は Parquet ファイル内のメタデータを自動的に活用し、不要なファイルや行グループ全体をスキップします。例えば、Parquet 形式の Hugging Face Dataset にフィルター（述語）を適用した場合や、カラムのサブセット（射影）を選択した場合、Dask は Parquet ファイルのメタデータを読み取り、ダウンロードせずに不要な部分を破棄します。

Dask DataFrame APIの再実装によるクエリ最適化のサポートが可能となり、Daskがより高速かつ堅牢になりました。

例えば、FineWeb-Eduのこのサブセットには多くのParquetファイルが含まれています。CCダンプの最近のテキストのみを保持するようにデータセットをフィルタリングできる場合、Daskはほとんどのファイルをスキップし、フィルタに一致するデータのみをダウンロードします。

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 は、計算に必要な列のみを読み取り、それ以外の列をスキップします。 例えば、コードの後半で列を削除する場合、パイプラインの初期段階でその列が不要であれば、読み込みに手間をかけません。 これは、列のサブセットを操作したい場合や、分析を行う場合に便利です。

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

クライアント

daskのほとんどの機能は、クラスタまたはローカルのClientによって並列計算を起動するように最適化されています:

import dask.dataframe as dd
from distributed import Client

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

ローカルでの使用において、Clientはデフォルトでマルチプロセッシングを使用するDaskのLocalClusterを使用します。LocalClusterのマルチプロセッシングを手動で設定するには、

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

Client を使用せずにローカルでデフォルトのスレッドスケジューラを使用する場合、特定の操作後に DataFrame が遅くなる可能性があることに注意してください（詳細はこちら）。

ローカルまたはクラウドクラスタのセットアップに関する詳細は、Daskのデプロイメントドキュメントを参照してください。