Dask

Dask là thư viện tính toán song song và phân tán, giúp mở rộng hệ sinh thái Python và PyData hiện có.

Cụ thể, chúng ta có thể sử dụng Dask DataFrame để mở rộng quy mô các quy trình làm việc của pandas. Dask DataFrame song song hóa pandas để xử lý dữ liệu bảng lớn. Nó mô phỏng sát API của pandas, giúp việc chuyển từ kiểm thử trên một tập dữ liệu đơn lẻ sang xử lý toàn bộ tập dữ liệu trở nên đơn giản. Dask đặc biệt hiệu quả với Parquet, định dạng mặc định trên Hugging Face Datasets, vì nó hỗ trợ các kiểu dữ liệu phong phú, lọc cột hiệu quả và nén dữ liệu.

Một trường hợp sử dụng thực tế tốt cho Dask là chạy xử lý dữ liệu hoặc suy luận mô hình trên một tập dữ liệu theo cách phân tán. Xem, ví dụ, bài đăng blog tuyệt vời của Coiled về Scaling AI-Based Data Processing with Hugging Face + Dask.

Đọc và Ghi

Vì Dask sử dụng fsspec để đọc và ghi dữ liệu từ xa, bạn có thể sử dụng các đường dẫn Hugging Face (hf://) để đọc và ghi dữ liệu trên Hub.

Trước tiên, bạn cần Đăng nhập bằng tài khoản Hugging Face, ví dụ bằng cách sử dụng:

hf auth login

Sau đó, bạn có thể Tạo kho dữ liệu, ví dụ bằng cách sử dụng:

from huggingface_hub import HfApi

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

Cuối cùng, bạn có thể sử dụng đường dẫn Hugging Face trong Dask. Dask DataFrame hỗ trợ ghi phân tán sang Parquet trên Hugging Face, trong đó sử dụng các commit để theo dõi các thay đổi của 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")

Vì điều này tạo ra một commit cho mỗi tệp, nên khuyến nghị là nén lịch sử sau khi tải lên:

from huggingface_hub import HfApi

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

Điều này tạo ra một kho dữ liệu username/my_dataset chứa tập dữ liệu Dask của bạn ở định dạng Parquet. Bạn có thể tải lại nó sau:

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

Để biết thêm thông tin về các đường dẫn của Hugging Face và cách chúng được triển khai, vui lòng tham khảo tài liệu của thư viện client về HfFileSystem.

Xử lý dữ liệu

Để xử lý một tập dữ liệu song song bằng Dask, bạn có thể trước tiên định nghĩa hàm xử lý dữ liệu của mình cho một DataFrame hoặc Series của pandas, sau đó sử dụng hàm map_partitions của Dask để áp dụng hàm này lên tất cả các phân vùng của tập dữ liệu một cách song song:

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

hoặc một hàm tương tự sử dụng các phương thức chuỗi của pandas (nhanh hơn):

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

Trong pandas, bạn có thể sử dụng hàm này trên một cột văn bản:

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

Và trong Dask, bạn có thể chạy hàm này trên mọi phân vùng:

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

Lưu ý rằng bạn cũng cần cung cấp meta, đây là kiểu của pandas Series hoặc DataFrame trong kết quả đầu ra của hàm của bạn. Điều này là cần thiết vì Dask DataFrame sử dụng API lazy. Vì Dask chỉ chạy xử lý dữ liệu khi .compute() được gọi, nên nó cần tham số meta để biết kiểu của cột mới trong thời gian chờ.

Đẩy xuống tiên đề và phép chiếu

Khi đọc dữ liệu Parquet từ Hugging Face, Dask tự động tận dụng metadata trong các tệp Parquet để bỏ qua toàn bộ tệp hoặc nhóm hàng (row group) nếu chúng không cần thiết. Ví dụ, nếu bạn áp dụng một bộ lọc (predicate) lên một Dataset Hugging Face ở định dạng Parquet hoặc nếu bạn chọn một tập con các cột (projection), Dask sẽ đọc metadata của các tệp Parquet để loại bỏ các phần không cần thiết mà không tải chúng xuống.

Điều này có thể thực hiện được nhờ sự tái triển khai API Dask DataFrame để hỗ trợ tối ưu hóa truy vấn, giúp Dask nhanh hơn và ổn định hơn.

Ví dụ, tập con này của FineWeb-Edu chứa nhiều tệp Parquet. Nếu bạn có thể lọc bộ dữ liệu để giữ lại văn bản từ các bản dump CC gần đây, Dask sẽ bỏ qua hầu hết các tệp và chỉ tải xuống dữ liệu khớp với bộ lọc:

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 cũng chỉ đọc các cột cần thiết cho phép tính của bạn và bỏ qua phần còn lại. Ví dụ, nếu bạn xóa một cột muộn trong mã của mình, nó sẽ không bận tâm đến việc tải nó sớm trong pipeline nếu không cần thiết. Điều này hữu ích khi bạn muốn thao tác trên một tập con các cột hoặc cho phân tích:

# 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

Hầu hết các tính năng trong dask được tối ưu hóa cho một cụm hoặc một Client cục bộ để khởi chạy các phép tính song song:

import dask.dataframe as dd
from distributed import Client

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

Để sử dụng cục bộ, Client sử dụng Dask LocalCluster với multiprocessing theo mặc định. Bạn có thể cấu hình thủ công multiprocessing của LocalCluster bằng

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

Lưu ý rằng nếu bạn sử dụng trình lập lịch thread mặc định cục bộ mà không có Client, một DataFrame có thể trở nên chậm hơn sau một số thao tác (chi tiết tại đây).

Tìm hiểu thêm thông tin về việc thiết lập cụm cục bộ hoặc cụm trên đám mây trong tài liệu triển khai Dask.