Hugging Face Datasets 라이브러리: 데이터 로드 및 처리

1. 개요

NLP 모델이나 LLM을 fine-tune할 때, 그 전에 항상 필요한 단계가 있습니다: 올바른 형식으로 데이터를 준비하고, 충분히 정제한 후, 학습 파이프라인에 로드하는 것입니다. Hugging Face의 datasets 라이브러리가 정확히 이 문제를 해결합니다. Hub에서 데이터셋을 로드하고, map/filter로 배치 처리하며, 반복 읽기에 최적화된 형식으로 저장할 수 있는 통합 API를 제공합니다.

이 글을 다 읽으면, Hugging Face Hub에서 데이터셋을 로드할 수 있고, Dataset/DatasetDict 구조를 이해하고, tokenizer로 텍스트를 전처리하고, streaming으로 대용량 데이터셋을 처리하며, 데이터 처리 파이프라인을 테스트하는 코드를 작성할 수 있을 것입니다.

2. 설치

pip install datasets transformers

datasets 라이브러리는 transformers에 하드 의존성이 없지만, 섹션 6의 tokenize 부분에서 필요합니다.

import datasets
print(datasets.__version__)

3.0.1

3. Hugging Face Hub에서 데이터셋 로드

3.1. load_dataset 기본 사용법

from datasets import load_dataset

reviews = load_dataset("imdb")
print(reviews)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 25000
    })
    unsupervised: Dataset({
        features: ['text', 'label'],
        num_rows: 50000
    })
})

load_dataset은 Hub에서 imdb 데이터셋의 스크립트/메타데이터를 로드하고, 압축된 파일을 로컬에 캐시(~/.cache/huggingface/datasets)한 다음, DatasetDict를 반환합니다. 이는 split(train, test, unsupervised)에 따라 정렬된 Dataset들의 딕셔너리입니다.

3.2. 특정 split만 로드하기

train_reviews = load_dataset("imdb", split="train")
print(type(train_reviews), len(train_reviews))

<class 'datasets.arrow_dataset.Dataset'> 25000

split="train"을 전달하면 DatasetDict 대신 Dataset을 직접 반환합니다. 한 가지 데이터셋만 필요할 때 편리합니다.

주의: Hub의 데이터셋이 config를 요구하는 경우(glue 등), config 이름을 추가로 전달해야 합니다: load_dataset("glue", "mrpc"). 이 파라미터가 없으면 config 선택을 요청하는 오류가 발생합니다.

4. Dataset 구조 탐색

sample = train_reviews[0]
print(sample.keys())
print(sample["label"], sample["text"][:60])

dict_keys(['text', 'label'])
0 I rented I AM CURIOUS-YELLOW from my video store because of...

인덱스로 접근하면 일반적인 Python 딕셔너리를 반환합니다. Dataset은 list들의 딕셔너리를 반환하는 slicing도 지원합니다:

batch = train_reviews[0:3]
print(batch["label"])

[0, 0, 0]

features를 사용해 스키마를 확인하세요:

print(train_reviews.features)

{'text': Value(dtype='string', id=None),
 'label': ClassLabel(names=['neg', 'pos'], id=None)}

ClassLabel은 0이 neg 레이블에 해당하고 1이 pos임을 알려줍니다. 레이블 이름으로 역매핑할 때 유용합니다.

5. map, filter, select로 데이터 처리

5.1. filter로 행 필터링하기

positive_only = train_reviews.filter(lambda row: row["label"] == 1)
print(len(positive_only))

12500

filter는 각 행에 대해 True/False를 반환하는 함수를 받아서 전체 데이터셋을 순회하고 조건을 만족하는 행만 유지합니다.

5.2. map으로 데이터 변환하기

def add_length(row):
    row["text_length"] = len(row["text"])
    return row

train_with_length = train_reviews.map(add_length)
print(train_with_length[0]["text_length"])

1640

map은 각 행에 함수를 적용하여 새로운 필드를 포함하는 딕셔너리를 반환합니다(또는 기존 필드를 업데이트). 결과는 원본 데이터셋에 병합됩니다.

5.3. map을 배치로 처리해 속도 높이기

def add_length_batched(batch):
    batch["text_length"] = [len(t) for t in batch["text"]]
    return batch

train_fast = train_reviews.map(add_length_batched, batched=True, batch_size=1000)

batched=True를 사용하면 함수가 각 행 대신 배치 딕셔너리(각 필드는 리스트) 전체를 받습니다. 이는 Python 함수 호출을 25,000번 피하고 대용량 데이터셋에서 훨씬 빠릅니다.

5.4. select로 부분집합 추출하기

small_train = train_reviews.select(range(100))
print(len(small_train))

100

전체 데이터셋을 실행하기 전에 부분집합에서 파이프라인을 빠르게 테스트하려 할 때 사용합니다.

6. tokenizer로 텍스트 전처리하기

가장 일반적인 실제 사용 사례: 모델에 입력하기 위해 텍스트를 tokenize합니다.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_batch(batch):
    return tokenizer(batch["text"], truncation=True, max_length=256)

tokenized = train_reviews.map(tokenize_batch, batched=True)
print(tokenized.column_names)

['text', 'label', 'input_ids', 'attention_mask']

tokenize_batch는 배치를 받아 input_ids, attention_mask를 포함하는 딕셔너리를 반환합니다. map은 이 새 필드들을 자동으로 데이터셋에 병합하고 기존 필드를 유지합니다. 불필요한 필드를 제거하려면(예: RAM 절약을 위해 원본 text), remove_columns을 사용하세요:

tokenized = train_reviews.map(
    tokenize_batch, batched=True, remove_columns=["text"]
)
print(tokenized.column_names)

['label', 'input_ids', 'attention_mask']

7. split 분할, shuffle, shard

7.1. train/test를 직접 분할하기

데이터셋에 split이 하나만 있다면 train_test_split으로 직접 분할할 수 있습니다:

split_result = train_reviews.train_test_split(test_size=0.2, seed=42)
print(split_result)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 20000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 5000
    })
})

seed=42는 여러 번의 실행 간에 분할 결과를 재현 가능하게 합니다.

7.2. Shuffle과 shard

shuffled = train_reviews.shuffle(seed=42)
part = shuffled.shard(num_shards=4, index=0)
print(len(part))

6250

shard는 데이터셋을 num_shards개의 동일한 부분으로 분할합니다. 분산 학습을 위해 여러 worker에 데이터를 배포할 때 유용합니다.

8. 대용량 데이터셋을 위한 Streaming

데이터셋이 너무 커서 디스크/RAM에 모두 로드할 수 없을 때 streaming=True를 사용하세요:

streamed = load_dataset("imdb", split="train", streaming=True)
first_three = list(streamed.take(3))
print(len(first_three))

3

streaming을 사용하면 datasets는 전체 파일을 로드하지 않고 iterator를 통해 각 레코드를 순차적으로 읽습니다. len()이나 직접 인덱싱은 사용할 수 없지만, lazy 방식으로 map/filter는 여전히 사용 가능합니다(iterate할 때만 적용됨).

주의: Streaming은 pretraining 세트와 같은 수백 GB의 데이터셋에 적합합니다. 작은 데이터셋(수백 MB)의 경우, 캐시를 활용할 수 있고 랜덤 액세스가 지원되므로 일반 로드(streaming=False)가 더 빠릅니다.

9. 데이터셋 저장 및 다시 로드하기

처리가 완료되면 결과를 저장하여 매번 map을 다시 실행할 필요가 없도록 합니다:

tokenized.save_to_disk("data/imdb_tokenized")

from datasets import load_from_disk

reloaded = load_from_disk("data/imdb_tokenized")
print(len(reloaded))

25000

save_to_disk는 데이터셋을 Arrow 형식으로 저장합니다. 원본에서 다시 parse할 필요가 없으므로 다시 읽기가 매우 빠릅니다. pandas와 상호작용해야 하면 직접 변환할 수 있습니다:

df = tokenized.select(range(5)).to_pandas()
print(df.columns.tolist())

['label', 'input_ids', 'attention_mask']

10. 데이터 처리 파이프라인 테스트 작성하기

데이터 처리 파이프라인은 코드를 수정할 때 map/filter 로직이 몰래 깨지지 않도록 테스트를 포함해야 합니다. 다음은 pytest를 사용한 테스트입니다.

import pytest
from datasets import Dataset


def add_length(row):
    row["text_length"] = len(row["text"])
    return row


def test_add_length_computes_correct_value():
    data = Dataset.from_dict({"text": ["hello", "hi there"]})

    result = data.map(add_length)

    assert result[0]["text_length"] == 5
    assert result[1]["text_length"] == 8


def test_filter_keeps_only_positive_labels():
    data = Dataset.from_dict({"label": [0, 1, 1, 0]})

    positive = data.filter(lambda row: row["label"] == 1)

    assert len(positive) == 2
    assert all(row["label"] == 1 for row in positive)

Dataset.from_dict는 메모리에 작은 데이터셋을 생성하며, Hub에서 아무것도 로드할 필요가 없습니다. unit test에 매우 적합합니다. 위의 두 테스트는 인터넷의 실제 데이터셋에 의존하지 않고 map(새 필드 생성)과 filter(올바른 조건으로 필터링)의 올바른 동작을 검증합니다.

pytest test_dataset_pipeline.py -v

test_dataset_pipeline.py::test_add_length_computes_correct_value PASSED
test_dataset_pipeline.py::test_filter_keeps_only_positive_labels PASSED

11. 일반적인 오류와 사용하지 말아야 할 경우

• 대용량 데이터셋 처리 시 batched=True 빠뜨리기: batch 없이 map을 실행하면 각 행마다 Python 함수를 호출하므로 느립니다. 수백만 행이 있을 때, 로직이 벡터화를 허용한다면 항상 batched=True를 우선시합니다.
• streaming 데이터셋에서 랜덤 인덱스 사용하기: streamed[5]는 오류를 발생시킵니다. streaming 데이터셋은 iterator이며 랜덤 액세스를 지원하지 않습니다. .take(), .skip() 또는 순차 iterate만 사용하세요.
• shuffle/train_test_split 시 seed 미설정하기: 실행 간 결과를 재현할 수 없으므로 실험 간 메트릭을 비교할 때 디버깅이 어렵습니다.
• 제어되지 않는 캐시 팽창: 다른 함수로 map을 실행할 때마다 datasets이 새로운 캐시 파일을 씁니다. 디스크 용량이 적으면 load_dataset(..., keep_in_memory=True)를 사용하거나 ~/.cache/huggingface/datasets를 정기적으로 삭제하세요.
• 복잡한 관계형 데이터에 datasets 사용하지 않기 (여러 테이블 join, 트랜잭션 필요). 이는 ML에 사용되는 평면형 표/텍스트 데이터에 최적화된 라이브러리이며, 데이터베이스의 대체품이 아닙니다.

12. 결론

datasets 라이브러리는 ’규모 큰 Hub에서 원본 데이터셋 로드’에서 ’토큰화된 학습 준비 완료 데이터셋’까지 load_dataset, map, filter 명령어 몇 개로만 이동할 수 있는 간단한 API를 제공합니다. 작은 데이터셋의 경우 정상적으로 로드하고 Arrow 캐시를 활용합니다. 거대한 데이터셋의 경우 streaming=True로 전환하여 전체를 디스크에 로드할 필요가 없습니다. 실제 프로젝트용 데이터 처리 파이프라인을 구축할 때 Dataset.from_dict를 사용하여 map/filter 함수의 테스트를 작성하는 것을 잊지 마세요. 이는 전처리 로직이 의도하지 않게 변경되었을 때 조기에 발견하는 데 도움이 됩니다.