데이터셋 스트리밍

데이터셋 스트리밍

Hugging Face에서 파일을 점진적으로 읽거나 일부만 읽는 스트리밍은 특히 큰 파일을 다룰 때 아주 유용해요.

출처: 문서

본문

통합 라이브러리

Hub의 데이터셋이 Hugging Face에서 스트리밍을 허용하는 지원 라이브러리와 호환된다면, 데이터셋 스트리밍은 몇 줄이면 끝나요. 데이터셋 접근 정보는 데이터셋 페이지의 "Use this dataset" 버튼을 클릭하면 볼 수 있고, knkarthik/samsum은 아래 datasets 라이브러리 사용법을 보여줘요.

Hugging Face 클라이언트 라이브러리 사용

huggingface_hub 라이브러리로 저장소의 파일을 생성·삭제·접근할 수 있어요. 예를 들어 Python에서 allenai/c4 데이터셋을 스트리밍하려면 라이브러리(최신 버전 권장)를 설치하고 다음 코드를 실행하면 돼요.

pip install -U huggingface_hub

from huggingface_hub import hffs

repo_id = "allenai/c4"
path_in_repo = "en/c4-train.00000-of-01024.json.gz"

# Stream the file
with hffs.open(f"datasets/{repo_id}/{path_in_repo}", "r", compression="gzip") as f:
    print(f.readline())  # read only the first line
    # {"text":"Beginners BBQ Class Taking Place in Missoula!...}

자세한 내용은 HfFileSystem 문서를 참고하세요.

이를 자신의 라이브러리에 통합할 수도 있어요! 예를 들어 Pandas로 CSV 데이터셋을 배치 단위로 빠르게 스트리밍할 수 있어요.

import pandas as pd

repo_id = "YOUR_REPO_ID"
path_in_repo = "data.csv"

batch_size = 5

# Stream the file
with hffs.open(f"datasets/{repo_id}/{path_in_repo}") as f:
    for df in pd.read_csv(f, iterator=True, chunksize=batch_size):  # read 5 lines at a time
        print(len(df))  # 5

스트리밍은 특히 Hugging Face의 큰 파일을 점진적으로 읽거나 아주 일부만 읽을 때 유용해요. 예를 들어 tarfile은 TAR 아카이브의 파일을, zipfile은 ZIP 아카이브의 파일을, pyarrow는 Parquet 파일의 row group에 접근할 수 있어요.

[!TIP] Rust에는 OpenDAL에 동등한 파일시스템 구현이 있어요.

cURL 사용

Hub의 모든 파일은 HTTP로 제공되므로 cURL로 스트리밍할 수 있어요.

>>> curl -L https://huggingface.co/datasets/fka/awesome-chatgpt-prompts/resolve/main/prompts.csv | head -n 5
"act","prompt"
"An Ethereum Developer","Imagine you are an experienced Ethereum developer tasked with creating...
"SEO Prompt","Using WebPilot, create an outline for an article that will be 2,000 words on the ...
"Linux Terminal","I want you to act as a linux terminal. I will type commands and you will repl...
"English Translator and Improver","I want you to act as an English translator, spelling correct...

범위 요청(range request)으로 파일의 특정 부분만 접근할 수 있어요.

>>> curl -r 40-88 -L https://huggingface.co/datasets/fka/awesome-chatgpt-prompts/resolve/main/prompts.csv
Imagine you are an experienced Ethereum developer

액세스 토큰으로 개인 저장소에서 스트리밍하세요.

>>> export HF_TOKEN=hf_xxx
>>> curl -H "Authorization: Bearer ***" -L https://huggingface.co/...

Parquet 스트리밍

Parquet은 AI 데이터셋에 아주 좋은 형식이에요. 좋은 압축률, 효율적인 처리와 프로젝션을 위한 컬럼 구조, 빠른 필터링을 위한 다단계 메타데이터를 제공하고, 모든 규모의 데이터셋에 적합해요.

Parquet 파일은 보통 각각 약 100MB인 row group으로 나뉘어요. 그래서 데이터 로더와 데이터 처리 프레임워크가 row group을 순회하며 데이터를 점진적으로 스트리밍할 수 있어요.

row group 안에는 개별 컬럼이 있고, 컬럼은 페이지로 나뉘어요. 페이지는 실제 데이터를 담는 약 1MB의 압축 블록이에요.

Row Group 스트리밍

PyArrow로 Hugging Face의 Parquet 파일에서 row group을 스트리밍하세요.

import pyarrow.parquet as pq

repo_id = "HuggingFaceFW/finewiki"
path_in_repo = "data/enwiki/000_00000.parquet"

# Stream the Parquet file row group per row group
with pq.ParquetFile(f"hf://datasets/{repo_id}/{path_in_repo}") as pf:
    for row_group_idx in range(pf.num_row_groups):
        row_group_table = pf.read_row_group(row_group_idx)
        df = row_group_table.to_pandas()

[!TIP] PyArrow는 hf:// 경로를 기본 지원하고 HfFileSystem을 자동으로 사용해요.

자세한 내용은 PyArrow 문서에서 확인할 수 있어요.

효율적인 랜덤 접근

Row group은 다시 컬럼으로, 컬럼은 페이지로 나뉘어요. 페이지는 보통 약 1MB이고 압축이 적용되는 지점이라 Parquet에서 데이터의 최소 단위예요. 페이지에 접근하면 전체 row group을 로드하지 않고도 특정 행을 로드할 수 있고, Parquet 파일에 페이지 인덱스가 있으면 가능해요. 하지만 모든 Parquet 프레임워크가 페이지 수준 읽기를 지원하는 것은 아니에요. 예를 들어 PyArrow는 지원하지 않지만 Rust의 parquet 크레이트는 지원해요.

use std::sync::Arc;
use object_store::path::Path;
use object_store_opendal::OpendalStore;
use opendal::services::Huggingface;
use opendal::Operator;
use parquet::arrow::async_reader::ParquetObjectReader;
use parquet::arrow::ParquetRecordBatchStreamBuilder;
use futures::TryStreamExt;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let repo_id = "HuggingFaceFW/finewiki";
    let path_in_repo = Path::from("data/enwiki/000_00000.parquet");
    let offset = 0;
    let limit = 10;

    let builder = Huggingface::default().repo_type("dataset").repo_id(repo_id);
    let operator = Operator::new(builder)?.finish();
    let store = Arc::new(OpendalStore::new(operator));
    let reader = ParquetObjectReader::new(store, path_in_repo.clone());
    let batch_stream =
        ParquetRecordBatchStreamBuilder::new(reader).await?
            .with_offset(offset as usize)
            .with_limit(limit as usize)
            .build()?;
    let results = batch_stream.try_collect::<Vec<_>>().await?;
    println!("Read {} batches", results.len());
    Ok(())
}

[!TIP] Rust에서는 Python의 HfFileSystem과 동등한 OpenDAL의 Huggingface 서비스를 사용해요.

PyArrow에서 write_page_index=True를 넘기면 효율적인 랜덤 접근을 지원하는 페이지 인덱스가 포함돼요. 특히 Parquet 컬럼에 "offset_index_offset"과 "offset_index_length"가 추가되고, 이는 Hugging Face의 Parquet 메타데이터 뷰어에서 확인할 수 있어요. 페이지 인덱스는 Hugging Face Dataset Viewer도 빨라지게 하고, row group 크기 제한 없이 데이터를 보여줄 수 있게 해요.

더 알아보기 (Learn more)

huggingface_hub의 hffs 또는 cURL의 범위 요청으로 큰 파일을 점진적으로 스트리밍할 수 있어요. Parquet은 row group 단위로 PyArrow가 hf:// 경로를 자동 지원하니, 큰 데이터셋은 row group을 순회하며 읽는 걸 권장해요. 페이지 인덱스(write_page_index=True)를 쓰면 특정 행만 효율적으로 읽을 수 있어요.