Datasets documentation
Gunakan dengan Polars
Gunakan dengan Polars
Dokumen ini adalah pengantar singkat tentang penggunaan datasets dengan Polars, dengan fokus khusus pada cara memproses dataset menggunakan fungsi Polars, serta cara mengonversi dataset ke Polars atau dari Polars.
Hal ini sangat berguna karena memungkinkan operasi zero-copy yang cepat, karena baik datasets maupun Polars menggunakan Arrow di baliknya.
Format dataset
Secara default, dataset mengembalikan objek Python biasa: bilangan bulat, bilangan desimal, string, daftar, dan sebagainya.
Untuk mendapatkan Polars DataFrames atau Series sebagai gantinya, Anda dapat mengatur format dataset menjadi polars menggunakan 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"
]Hal ini juga berlaku untuk objek IterableDataset yang diperoleh, misalnya, menggunakan 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 │
└───────┴───────┘Proses data
Fungsi Polars umumnya lebih cepat daripada fungsi python yang ditulis manual, sehingga merupakan pilihan yang baik untuk mengoptimalkan pemrosesan data. Anda dapat menggunakan fungsi Polars untuk memproses dataset di Dataset.map() atau 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 │
└───────┴───────┴───────┘Kami menggunakan batched=True karena memproses data dalam bentuk batch di Polars lebih cepat daripada baris per baris. Anda juga dapat menggunakan batch_size= di map() untuk mengatur ukuran setiap df.
Hal ini juga berlaku untuk IterableDataset.map() dan IterableDataset.filter().
Contoh: ekstraksi data
Banyak fungsi tersedia di Polars untuk semua tipe data: string, float, integer, dan lain-lain. Anda dapat menemukan daftar lengkapnya di sini. Fungsi-fungsi tersebut ditulis dalam Rust dan dijalankan pada batch data, yang memungkinkan pemrosesan data yang cepat.
Berikut adalah contoh yang menunjukkan peningkatan kecepatan 5x dengan menggunakan Polars alih-alih fungsi python biasa untuk mengekstrak solusi dari dataset penalaran 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: 2sImport atau Export dari Polars
Untuk mengimpor data dari Polars, Anda dapat menggunakan Dataset.from_polars():
ds = Dataset.from_polars(df)
Dan Anda dapat menggunakan Dataset.to_polars() untuk mengekspor Dataset ke Polars DataFrame:
df = Dataset.to_polars(ds)