Datasets documentation
Sử dụng với Polars
Sử dụng với Polars
Tài liệu này là phần giới thiệu nhanh về việc sử dụng datasets với Polars, tập trung đặc biệt vào cách xử lý datasets bằng các hàm của Polars, cũng như cách chuyển đổi một dataset sang Polars hoặc từ Polars.
Điều này đặc biệt hữu ích vì nó cho phép các thao tác zero-copy nhanh chóng, do cả datasets và Polars đều sử dụng Arrow ở bên trong.
Định dạng tập dữ liệu
Mặc định, các tập dữ liệu trả về các đối tượng Python thông thường: số nguyên, số thực, chuỗi, danh sách, v.v.
Để nhận các DataFrame hoặc Series của Polars thay thế, bạn có thể đặt định dạng của dataset thành polars bằng cách sử dụng Dataset.with_format():
>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("polars")
>>> ds[0] # pl.DataFrame
shape: (1, 2)
┌───────┬───────┐
│ col_0 ┆ col_1 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═══════╪═══════╡
│ a ┆ 0.0 │
└───────┴───────┘
>>> ds[:2] # pl.DataFrame
shape: (2, 2)
┌───────┬───────┐
│ col_0 ┆ col_1 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═══════╪═══════╡
│ a ┆ 0.0 │
│ b ┆ 0.0 │
└───────┴───────┘
>>> ds["data"] # pl.Series
shape: (4,)
Series: 'col_0' [str]
[
"a"
"b"
"c"
"d"
]Điều này cũng hoạt động với các đối tượng IterableDataset được lấy, ví dụ bằng cách sử dụng load_dataset(..., streaming=True):
>>> ds = ds.with_format("polars")
>>> for df in ds.iter(batch_size=2):
... print(df)
... break
shape: (2, 2)
┌───────┬───────┐
│ col_0 ┆ col_1 │
│ --- ┆ --- │
│ str ┆ f64 │
╞═══════╪═══════╡
│ a ┆ 0.0 │
│ b ┆ 0.0 │
└───────┴───────┘Xử lý dữ liệu
Các hàm Polars thường nhanh hơn so với các hàm python viết tay thông thường, do đó chúng là một lựa chọn tốt để tối ưu hóa xử lý dữ liệu. Bạn có thể sử dụng các hàm Polars để xử lý một tập dữ liệu trong Dataset.map() hoặc Dataset.filter():
>>> import polars as pl
>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("polars")
>>> ds = ds.map(lambda df: df.with_columns(pl.col("col_1").add(1).alias("col_2")), batched=True)
>>> ds[:2]
shape: (2, 3)
┌───────┬───────┬───────┐
│ col_0 ┆ col_1 ┆ col_2 │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═══════╪═══════╪═══════╡
│ a ┆ 0.0 ┆ 1.0 │
│ b ┆ 0.0 ┆ 1.0 │
└───────┴───────┴───────┘
>>> ds = ds.filter(lambda df: df["col_0"] == "b", batched=True)
>>> ds[0]
shape: (1, 3)
┌───────┬───────┬───────┐
│ col_0 ┆ col_1 ┆ col_2 │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 │
╞═══════╪═══════╪═══════╡
│ b ┆ 0.0 ┆ 1.0 │
└───────┴───────┴───────┘Chúng tôi sử dụng batched=True vì việc xử lý dữ liệu theo lô trong Polars nhanh hơn so với xử lý từng hàng. Bạn cũng có thể sử dụng batch_size= trong map() để đặt kích thước của mỗi df.
Điều này cũng hoạt động với IterableDataset.map() và IterableDataset.filter().
Ví dụ: trích xuất dữ liệu
Nhiều hàm có sẵn trong Polars và hỗ trợ mọi kiểu dữ liệu: chuỗi, số thực, số nguyên, v.v. Bạn có thể tìm thấy danh sách đầy đủ tại đây. Những hàm này được viết bằng Rust và chạy trên các lô dữ liệu, cho phép xử lý dữ liệu nhanh chóng.
Dưới đây là một ví dụ cho thấy tốc độ tăng gấp 5 lần khi sử dụng Polars thay vì một hàm python thông thường để trích xuất các lời giải từ một tập dữ liệu suy luận của LLM:
from datasets import load_dataset
ds = load_dataset("ServiceNow-AI/R1-Distill-SFT", "v0", split="train")
# Using a regular python function
pattern = re.compile("boxed\\{(.*)\\}")
result_ds = ds.map(lambda x: {"value_solution": m.group(1) if (m:=pattern.search(x["solution"])) else None})
# Time: 10s
# Using a Polars function
expr = pl.col("solution").str.extract("boxed\\{(.*)\\}").alias("value_solution")
result_ds = ds.with_format("polars").map(lambda df: df.with_columns(expr), batched=True)
# Time: 2sNhập hoặc Xuất từ Polars
Để nhập dữ liệu từ Polars, bạn có thể sử dụng Dataset.from_polars():
ds = Dataset.from_polars(df)
Và bạn có thể sử dụng Dataset.to_polars() để xuất một Dataset sang một DataFrame Polars:
df = Dataset.to_polars(ds)