Datasets documentation
Polars와 함께 사용
Polars와 함께 사용
이 문서는 datasets를 Polars와 함께 사용하는 방법에 대한 간단한 소개입니다. 특히 Polars 함수를 사용하여 데이터를 처리하는 방법과, Polars로 데이터를 변환하거나 Polars에서 데이터를 가져오는 방법에 중점을 둡니다.
datasets와 Polars 모두 내부적으로 Arrow를 사용하므로, 이는 빠른 제로-복사(zero-copy) 연산을 가능하게 해주어 특히 유용합니다.
데이터셋 형식
기본적으로 datasets는 정수, 실수, 문자열, 리스트 등 일반적인 Python 객체를 반환합니다.
To get Polars DataFrames or Series instead, you can set the format of the dataset to polars using 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"
]IterableDataset 객체, 예를 들어 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 │
└───────┴───────┘데이터 처리
Polars 함수는 일반적으로 직접 작성한 Python 함수보다 빠르므로 데이터 처리를 최적화하는 데 좋은 선택지입니다. Dataset.map() 또는 Dataset.filter()에서 Polars 함수를 사용해 데이터셋을 처리할 수 있습니다:
>>> 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 │
└───────┴───────┴───────┘Polars에서 데이터를 행별로 처리하는 것보다 배치 단위로 처리하는 것이 더 빠르기 때문에 batched=True를 사용합니다. 또한 map()에서 batch_size=를 사용하여 각 df의 크기를 설정할 수도 있습니다.
이것은 IterableDataset.map()과 IterableDataset.filter()에도 적용됩니다.
예: 데이터 추출
Polars에는 문자열, 실수, 정수 등 모든 데이터 타입에 사용할 수 있는 다양한 함수가 있습니다. 전체 목록은 여기에서 확인할 수 있습니다. 이러한 함수는 Rust로 작성되어 데이터 배치(batch) 단위로 실행되며, 이를 통해 빠른 데이터 처리가 가능합니다.
LLM 추론 데이터셋에서 솔루션을 추출할 때 일반 python 함수 대신 Polars를 사용하여 5배의 속도 향상을 보여주는 예시입니다:
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: 2sPolars에서 가져오기 또는 내보내기
Polars에서 데이터를 가져오려면 Dataset.from_polars()를 사용할 수 있습니다:
ds = Dataset.from_polars(df)
Dataset.to_polars()를 사용하여 Dataset을 Polars DataFrame으로 내보낼 수 있습니다:
df = Dataset.to_polars(ds)