Datasets documentation

스트림

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v4.8.4).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

스트림

Dataset streaming은 데이터셋을 다운로드하지 않고도 작업할 수 있게 해줍니다. 데이터는 데이터셋을 반복할 때 스트리밍됩니다. 다음과 같은 경우 특히 유용합니다:

  • 매우 큰 데이터셋의 다운로드를 기다리고 싶지 않을 것입니다.
  • 데이터셋 크기가 컴퓨터의 사용 가능한 디스크 공간을 초과했습니다.
  • 데이터셋의 일부 샘플만 빠르게 탐색하고 싶을 때

예를 들어, HuggingFaceFW/fineweb 데이터셋의 영어 분할은 45테라바이트이지만, 스트리밍을 사용하면 즉시 사용할 수 있습니다. 아래에 표시된 것처럼 load_dataset()에서 streaming=True를 설정하여 데이터셋을 스트리밍합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> print(next(iter(dataset)))
{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012\nWhen the first serious bout of tornadoes of 2012 blew through middle America in the middle of the night, they touched down in places hours from any AP bureau...', ...,
 'language_score': 0.9721424579620361, 'token_count': 717}

Dataset streaming을 사용하면 변환 과정 없이 로컬 파일로 구성된 dataset을 사용할 수 있습니다. 이 경우 dataset을 반복하는 동안 로컬 파일에서 데이터가 스트리밍됩니다. 다음과 같은 상황에서 특히 유용합니다:

  • 매우 큰 로컬 데이터셋이 Arrow로 변환되는 것을 기다리고 싶지 않을 것입니다.
  • 변환된 파일의 크기가 컴퓨터의 사용 가능한 디스크 공간을 초과합니다.
  • 데이터셋의 일부 샘플만 빠르게 탐색하고 싶을 때
  • Parquet 데이터셋에서 특정 열만 로드하거나 효율적으로 필터링하려면

예를 들어, oscar-corpus/OSCAR-2201와 같이 수백 개의 압축된 JSONL 파일로 구성된 로컬 데이터셋을 스트리밍하여 즉시 사용할 수 있습니다:

>>> from datasets import load_dataset
>>> data_files = {'train': 'path/to/OSCAR-2201/compressed/en_meta/*.jsonl.gz'}
>>> dataset = load_dataset('json', data_files=data_files, split='train', streaming=True)
>>> print(next(iter(dataset)))
{'id': 0, 'text': 'Founded in 2015, Golden Bees is a leading programmatic recruitment platform dedicated to employers, HR agencies and job boards. The company has developed unique HR-custom technologies and predictive algorithms to identify and attract the best candidates for a job opportunity.', ...

Parquet는 원하는 컬럼의 부분 집합만 스트리밍 및 로드하고 원하지 않는 컬럼을 무시할 수 있는 컬럼형 형식입니다. Parquet는 파일 및 row group 수준의 컬럼 통계와 같은 메타데이터를 저장하여 효율적인 필터링을 지원합니다. datasets.packaged_modules.parquet.ParquetConfig의 columns 및 filters 인수를 사용하여 Parquet 데이터셋을 스트리밍하고 컬럼을 선택하며 필터를 적용합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, columns=["url", "date"])
>>> print(next(iter(dataset)))
{'url': 'http://%20jwashington@ap.org/Content/Press-Release/2012/How-AP-reported-in-all-formats-from-tornado-stricken-regions', 'date': '2013-05-18T05:48:54Z'}
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, filters=[("language_score", ">=", 0.99)])
>>> print(next(iter(dataset)))
{'text': 'Everyone wishes for something. And lots of people believe they know how to make their wishes come true with magical thinking.\nWhat is it? "Magical thinking is a belief in forms of causation, with no known physical basis," said Professor Emily Pronin of Princeton...', ...,
 'language_score': 0.9900368452072144, 'token_count': 716}

스트리밍 모드로 데이터셋을 로드하면 Dataset 객체가 아닌 IterableDataset이라는 새로운 데이터셋 타입 인스턴스가 생성됩니다. 이 특수한 데이터셋 타입에는 아래에 표시된 자체 처리 메서드 세트가 있습니다.

An IterableDataset is useful for iterative jobs like training a model. You shouldn’t use a IterableDataset for jobs that require random access to examples because you have to iterate all over it using a for loop. Getting the last example in an iterable dataset would require you to iterate over all the previous examples. You can find more details in the Dataset vs. IterableDataset guide.

Column indexing

특정 열의 값을 반복하는 것이 편리한 경우가 있습니다. 다행히도 IterableDataset은 열 인덱싱을 지원합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> print(next(iter(dataset["text"])))
Beginners BBQ Class Taking Place in Missoula!...

Dataset에서 변환

기존 Dataset 객체가 있는 경우, to_iterable_dataset() 함수를 사용하여 IterableDataset으로 변환할 수 있습니다. 이는 load_dataset()에서 streaming=True 인수를 설정하는 것보다 실제로 더 빠릅니다. 데이터가 로컬 파일에서 스트리밍되기 때문입니다.

>>> from datasets import load_dataset

# faster 🐇
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset()

# slower 🐢
>>> iterable_dataset = load_dataset("ethz/food101", streaming=True)

The to_iterable_dataset() function supports sharding when the IterableDataset is instantiated. This is useful when working with big datasets, and you’d like to shuffle the dataset or to enable fast parallel loading with a PyTorch DataLoader.

>>> import torch
>>> from datasets import load_dataset

>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset(num_shards=64) # shard the dataset
>>> iterable_dataset = iterable_dataset.shuffle(buffer_size=10_000)  # shuffles the shards order and use a shuffle buffer when you start iterating
dataloader = torch.utils.data.DataLoader(iterable_dataset, num_workers=4)  # assigns 64 / 4 = 16 shards from the shuffled list of shards to each worker when you start iterating

셔플

일반적인 Dataset 객체와 마찬가지로, IterableDataset도 IterableDataset.shuffle()를 사용하여 섞을 수 있습니다.

buffer_size 인수는 무작위로 예시를 샘플링할 버퍼의 크기를 제어합니다. 데이터셋에 100만 개의 예시가 있고 buffer_size를 10,000으로 설정했다고 가정해 보겠습니다. IterableDataset.shuffle()는 버퍼의 첫 번째 10,000개 예시에서 무작위로 예시를 선택합니다. 버퍼에서 선택된 예시는 새로운 예시로 대체됩니다. 기본적으로 버퍼 크기는 1,000입니다.

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> shuffled_dataset = dataset.shuffle(seed=42, buffer_size=10_000)

IterableDataset.shuffle()는 버퍼를 채울 때 한 번에 최대 max_buffer_input_shards개의 샤드를 사용합니다(기본값은 10). 이를 통해 셔플링의 품질이 크게 향상됩니다. 또한 데이터셋이 여러 파일로 샤딩된 경우 샤드의 순서도 셔플링합니다.

파일이 하나이거나 소수인 Parquet 데이터셋을 스트리밍하는 경우, 셔플링 전에 reshard()를 사용하세요. 이는 데이터셋을 파일 단위가 아닌 row group 단위로 샤딩하여 셔플링 품질을 향상시키는 데 유용합니다.

재배치

때로는 각 epoch 이후에 dataset을 다시 섞고 싶을 수 있습니다. 이를 위해 각 epoch마다 다른 seed를 설정해야 합니다. epoch 사이에 IterableDataset.set_epoch()를 사용하여 dataset에 현재 epoch을 알려주세요.

시드는 사실상 initial seed + current epoch가 됩니다.

>>> for epoch in range(epochs):
...     shuffled_dataset.set_epoch(epoch)
...     for example in shuffled_dataset:
...         ...

데이터셋 분할

데이터셋을 두 가지 방식으로 나눌 수 있습니다:

>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> dataset_head = dataset.take(2)
>>> list(dataset_head)
[{'text': "How AP reported in all formats from tor...},
 {'text': 'Did you know you have two little yellow...}]
  • IterableDataset.skip()는 데이터셋의 첫 n개 예시를 생략하고 나머지 예시를 반환합니다:
>>> train_dataset = shuffled_dataset.skip(1000)

take와 skip은 샤드의 순서를 고정하기 때문에 shuffle의 향후 호출을 방지합니다. 데이터셋을 분할하기 전에 shuffle해야 합니다.

Shard

🤗 Datasets는 매우 큰 데이터셋을 미리 정의된 개수의 청크로 나누는 샤딩을 지원합니다. shard()에서 num_shards 매개변수를 지정하여 데이터셋을 몇 개의 샤드로 나눌지 결정합니다. 또한 index 매개변수를 통해 반환할 샤드를 지정해야 합니다.

예를 들어, amazon_polarity 데이터셋에는 4개의 샤드(이 경우 4개의 Parquet 파일)가 있습니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset("fancyzhx/amazon_polarity", split="train", streaming=True)
>>> print(dataset)
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 4
})

데이터셋을 두 개의 청크로 샤딩한 후, 첫 번째 청크에는 샤드가 2개만 있습니다:

>>> dataset.shard(num_shards=2, index=0)
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 2
})

데이터셋의 샤드 수를 늘리려면 IterableDataset.reshard()를 사용할 수 있습니다:

>>> dataset.reshard()
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 3600
})

리샤딩 메커니즘은 데이터셋 파일 형식에 따라 달라집니다. 예를 들어 Parquet의 경우, 샤드당 하나의 파일 대신 row group을 사용하여 리샤딩합니다. 각 형식에서 작동하는 방식은 IterableDataset.reshard()의 문서에서 확인하세요.

리샤딩 후에도 데이터셋의 dataset.num_shards==1인 경우, 대신 IterableDataset.skip()와 IterableDataset.take()를 사용하여 청킹해야 합니다.

Concatenate

같은 컬럼 타입을 공유하는 개별 데이터셋은 연결할 수 있습니다. concatenate_datasets()로 데이터셋을 연결합니다:

>>> from datasets import concatenate_datasets, load_dataset

>>> stories = load_dataset("ajibawa-2023/General-Stories-Collection", split="train", streaming=True)
>>> stories = stories.select_columns(["text"])  # only keep the 'text' column
IterableDataset({
    features: Unknown,
    num_shards: 10
})

>>> wiki = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
>>> wiki = wiki.select_columns(["text"])  # only keep the 'text' column
IterableDataset({
    features: ['text'],
    num_shards: 41
})

>>> bert_dataset = concatenate_datasets([stories, wiki])
IterableDataset({
    features: ['text'],
    num_shards: 51
})

concatenated 데이터셋의 샤드는 입력 데이터셋의 샤드입니다.

행 수가 동일한 데이터셋인 경우 axis=1을 설정하여 두 데이터셋을 수평으로 연결할 수도 있습니다:

>>> from datasets import IterableDataset
>>> stories_ids = IterableDataset.from_dict({"ids": list(range(num_stories))})
>>> stories_with_ids = concatenate_datasets([stories, stories_ids], axis=1)

이 경우, 입력 데이터셋에서 정렬되지 않은 샤드로 끝나지 않도록 하기 위해 연결된 데이터셋은 샤드가 1개만 있습니다.

교차

interleave_datasets()는 동일한 컬럼 유형을 가진 IterableDataset와 다른 데이터셋을 결합할 수 있습니다. 결합된 데이터셋은 원래 데이터셋들에서 번갈아 가며 예시를 반환합니다.

>>> from datasets import interleave_datasets
>>> es_dataset = load_dataset('allenai/c4', 'es', split='train', streaming=True)
IterableDataset({
    features: Unknown,
    num_shards: 2048
})
>>> fr_dataset = load_dataset('allenai/c4', 'fr', split='train', streaming=True)
IterableDataset({
    features: Unknown,
    num_shards: 2048
})

>>> multilingual_dataset = interleave_datasets([es_dataset, fr_dataset])
IterableDataset({
    features: ['text', 'timestamp', 'url'],
    num_shards: 2048
})
>>> list(multilingual_dataset.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
 {'text': 'Le sacre de philippe ier, 23 mai 1059 - Compte Rendu...'}]

각 원본 데이터셋에서 샘플링할 확률을 정의하여, 각 데이터셋이 어떻게 샘플링되고 결합되는지에 대한 제어를 강화합니다. probabilities 인자에 원하는 샘플링 확률을 설정하세요:

>>> multilingual_dataset_with_oversampling = interleave_datasets([es_dataset, fr_dataset], probabilities=[0.8, 0.2], seed=42)
>>> list(multilingual_dataset_with_oversampling.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
 {'text': 'Chevrolet Cavalier Usados en Bogota - Carros en Vent...'}]

최종 데이터셋의 약 80%는 es_dataset로, 20%는 fr_dataset로 구성됩니다.

stopping_strategy도 지정할 수 있습니다. 기본 전략인 first_exhausted는 서브샘플링 전략으로, 즉 어떤 데이터셋의 샘플이 먼저 소진되면 데이터셋 구성이 중단됩니다. stopping_strategy=all_exhausted를 지정하여 오버샘플링 전략을 실행할 수 있습니다. 이 경우, 모든 데이터셋의 모든 샘플이 최소 한 번씩 추가되면 데이터셋 구성이 중단됩니다. 실제로는 데이터셋이 소진되면 중단 조건이 충족될 때까지 해당 데이터셋의 처음부터 다시 시작합니다. 샘플링 확률이 지정되지 않은 경우, 새로운 데이터셋에는 max_length_datasets*nb_dataset samples이 포함됩니다. 모든 샘플이 정확히 한 번만 보이도록 보장하는 stopping_strategy=all_exhausted_without_replacement도 있습니다.

sharding를 사용하여 적절한 병렬성을 보장하기 위해, interleaved 데이터셋의 각 shard에는 모든 입력 데이터셋의 shard가 최소 1개 이상 포함되어 있습니다. 따라서 interleaved 데이터셋의 sharding 레벨은 입력 데이터셋 중 가장 낮은 sharding 레벨이 됩니다.

예를 들어, 입력 데이터셋의 샤드가 각각 32개, 48개, 128개라면, 인터리브된 데이터셋은 32 = min(32, 48, 128)개의 샤드를 가지며, 각 새로운 샤드는 첫 번째 데이터셋에서 1개, 두 번째 데이터셋에서 1~2개, 세 번째 데이터셋에서 4개의 샤드를 포함합니다.

rename, remove, cast

다음 메서드를 사용하면 데이터셋의 열을 수정할 수 있습니다. 이러한 메서드는 열 이름 변경 또는 열 제거, 열을 새로운 특징 집합으로 변경하는 데 유용합니다.

이름 변경

데이터셋의 컬럼 이름을 변경해야 할 때는 IterableDataset.rename_column()을 사용하세요. 원래 컬럼에 연결된 Features는 단순히 원본 컬럼을 제자리에서 대체하는 것이 아니라, 실제로 새 컬럼 이름 아래로 이동됩니다.

IterableDataset.rename_column()에 원본 컬럼 이름과 새 컬럼 이름을 제공합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.rename_column("text", "content")

제거

하나 이상의 컬럼을 제거해야 할 경우, IterableDataset.remove_columns()에 제거할 컬럼 이름을 전달합니다. 여러 컬럼을 제거하려면 컬럼 이름 목록을 제공합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.remove_columns('timestamp')

캐스팅

IterableDataset.cast()는 하나 이상의 컬럼의 feature type을 변경합니다. 이 메서드는 새로운 Features를 인수로 받습니다. 다음 샘플 코드는 ClassLabel와 Value의 feature type을 변경하는 방법을 보여줍니다.

>>> from datasets import load_dataset
>>> dataset = load_dataset('nyu-mll/glue', 'mrpc', split='train', streaming=True)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['not_equivalent', 'equivalent']),
'idx': Value('int32')}

>>> from datasets import ClassLabel, Value
>>> new_features = dataset.features.copy()
>>> new_features["label"] = ClassLabel(names=['negative', 'positive'])
>>> new_features["idx"] = Value('int64')
>>> dataset = dataset.cast(new_features)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['negative', 'positive']),
'idx': Value('int64')}

Casting은 원래 feature type과 새 feature type이 호환되는 경우에만 작동합니다. 예를 들어, 원래 컬럼에 1과 0만 포함된 경우, feature type이 Value('int32')인 컬럼을 Value('bool')로 cast할 수 있습니다.

IterableDataset.cast_column()을 사용하여 특정 열 하나의 feature type만 변경할 수 있습니다. 열 이름과 새로운 feature type을 인수로 전달하세요:

>>> dataset.features
{'audio': Audio(sampling_rate=44100, mono=True)}

>>> dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
>>> dataset.features
{'audio': Audio(sampling_rate=16000, mono=True)}

맵

일반 Dataset에 사용되는 Dataset.map() 함수와 유사하게, 🤗 Datasets에는 IterableDataset을 처리하기 위한 IterableDataset.map() 기능이 있습니다. IterableDataset.map()은 예제가 스트리밍되는 시점에 온더플라이(on-the-fly)로 처리를 적용합니다.

데이터셋의 각 예시에 대해 처리 함수를 독립적으로 또는 배치 단위로 적용할 수 있습니다. 이 함수는 새로운 행과 열을 생성할 수도 있습니다.

다음 예제는 IterableDataset를 토큰화하는 방법을 보여줍니다. 이 함수는 dict를 받아서 출력해야 합니다:

>>> def add_prefix(example):
...     example['text'] = 'My text: ' + example['text']
...     return example

다음으로, IterableDataset.map()를 사용하여 이 함수를 데이터셋에 적용합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> updated_dataset = dataset.map(add_prefix)
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...',
  'timestamp': '2019-04-25 12:57:54',
  'url': 'https://klyq.com/beginners-bbq-class-taking-place-in-missoula/'},
 {'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...',
  'timestamp': '2019-04-21 10:07:13',
  'url': 'https://forums.macrumors.com/threads/restore-from-larger-disk-to-smaller-disk.1311329/'},
 {'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...',
  'timestamp': '2019-04-25 10:40:23',
  'url': 'https://awishcometrue.com/Catalogs/Clearance/Tweens/V1960-Find-A-Way'}]

다른 예제를 살펴보겠습니다. 이번에는 IterableDataset.map()을 사용하여 컬럼을 제거합니다. 컬럼을 제거하면 해당 예제가 매핑된 함수에 제공된 후에야 실제로 제거됩니다. 이를 통해 매핑된 함수는 컬럼이 제거되기 전에 해당 컬럼의 내용을 사용할 수 있습니다.

IterableDataset.map()에서 remove_columns 인수를 사용하여 제거할 열을 지정합니다:

>>> updated_dataset = dataset.map(add_prefix, remove_columns=["timestamp", "url"])
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...'},
 {'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...'},
 {'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...'}]

배치 처리

IterableDataset.map()은 예시 배치로 작업하는 것도 지원합니다. batched=True를 설정하여 배치 단위로 작업을 수행합니다. 기본 배치 크기는 1000이지만, batch_size 인수를 통해 조정할 수 있습니다. 이를 통해 토큰화, 긴 문장을 더 짧은 청크로 분할, 데이터 증강 등 다양한 흥미로운 애플리케이션이 가능해집니다.

토큰화

>>> from datasets import load_dataset
>>> from transformers import AutoTokenizer
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
>>> def encode(examples):
...     return tokenizer(examples['text'], truncation=True, padding='max_length')
>>> dataset = dataset.map(encode, batched=True, remove_columns=["text", "timestamp", "url"])
>>> next(iter(dataset))
{'input_ids': [101, 4088, 16912, 22861, 4160, 2465, 2635, 2173, 1999, 3335, ..., 0, 0, 0],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 0, 0]}

See other examples of batch processing in the batched map processing documentation. They work the same for iterable datasets.

필터

Dataset.filter()를 사용하여 조건 함수에 따라 데이터셋의 행을 필터링할 수 있습니다. 지정된 조건과 일치하는 행을 반환합니다:

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', streaming=True, split='train')
>>> start_with_ar = dataset.filter(lambda example: example['text'].startswith('San Francisco'))
>>> next(iter(start_with_ar))
{'text': 'San Francisco 49ers cornerback Shawntae Spencer will miss the rest of the sea...}

with_indices=True를 설정하면 Dataset.filter()는 인덱스로도 필터링할 수 있습니다:

>>> even_dataset = dataset.filter(lambda example, idx: idx % 2 == 0, with_indices=True)
>>> list(even_dataset.take(3))
[{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012 Whe...},
 {'text': 'Car Wash For Clara! Now is your chance to help! 2 year old Clara Woodward has...},
 {'text': 'Log In Please enter your ECode to log in. Forgotten your eCode? If you create...}]

배치

batch 메서드는 IterableDataset을 배치의 반복 가능한 객체로 변환합니다. 이는 학습 루프에서 배치를 다루거나 배치 입력을 기대하는 프레임워크를 사용할 때 특히 유용합니다.

map 함수를 사용하여 데이터 배치에 함수를 적용할 때 사용할 수 있는 “Batch Processing” 옵션도 있으며, 이는 위의 Map 섹션에서 다루고 있습니다. 여기서 설명하는 batch 메서드는 이와 다르며, 데이터셋에서 배치를 생성하는 보다 직접적인 방법을 제공합니다.

batch 메서드를 다음과 같이 사용할 수 있습니다:

from datasets import load_dataset

# Load a dataset in streaming mode
dataset = load_dataset("some_dataset", split="train", streaming=True)

# Create batches of 32 samples
batched_dataset = dataset.batch(batch_size=32)

# Iterate over the batched dataset
for batch in batched_dataset:
    print(batch)
    break

이 예시에서 batched_dataset은 여전히 IterableDataset이지만, 이제 각 항목은 단일 샘플이 아니라 32개의 샘플로 구성된 배치입니다. 이러한 배치는 데이터셋을 반복하는 동안 실시간으로 수행되어 IterableDataset의 메모리 효율적인 특성을 유지합니다.

batch 메서드는 drop_last_batch 파라미터도 제공합니다. True로 설정하면, 마지막 배치의 크기가 지정된 batch_size보다 작을 경우 해당 배치를 삭제합니다. 이 기능은 하위 처리 과정에서 모든 배치가 동일한 크기를 가져야 하는 시나리오에서 유용할 수 있습니다:

batched_dataset = dataset.batch(batch_size=32, drop_last_batch=True)

학습 루프에서 스트리밍

IterableDataset는 학습 루프에 통합할 수 있습니다. 먼저 데이터셋을 섞습니다:

Pytorch
Hide Pytorch content
>>> seed, buffer_size = 42, 10_000
>>> dataset = dataset.shuffle(seed, buffer_size=buffer_size)

마지막으로, 간단한 학습 루프를 만들고 학습을 시작합니다:

>>> import torch
>>> from torch.utils.data import DataLoader
>>> from transformers import AutoModelForMaskedLM, DataCollatorForLanguageModeling
>>> from tqdm import tqdm
>>> dataset = dataset.with_format("torch")
>>> dataloader = DataLoader(dataset, collate_fn=DataCollatorForLanguageModeling(tokenizer))
>>> device = 'cuda' if torch.cuda.is_available() else 'cpu' 
>>> model = AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased")
>>> model.train().to(device)
>>> optimizer = torch.optim.AdamW(params=model.parameters(), lr=1e-5)
>>> for epoch in range(3):
...     dataset.set_epoch(epoch)
...     for i, batch in enumerate(tqdm(dataloader, total=5)):
...         if i == 5:
...             break
...         batch = {k: v.to(device) for k, v in batch.items()}
...         outputs = model(**batch)
...         loss = outputs[0]
...         loss.backward()
...         optimizer.step()
...         optimizer.zero_grad()
...         if i % 10 == 0:
...             print(f"loss: {loss}")

데이터셋 체크포인트를 저장하고 반복을 재개합니다

학습 루프가 중단된 경우, 중단된 지점부터 학습을 재개할 수 있습니다. 이를 위해 모델과 옵티마이저, 그리고 데이터 로더의 체크포인트를 저장할 수 있습니다.

Iterable datasets는 특정 예시 인덱스에 대한 랜덤 액세스를 제공하여 해당 위치에서 재개할 수 없으므로, 모델과 옵티마이저에서 수행할 수 있는 것과 유사하게 IterableDataset.state_dict()와 IterableDataset.load_state_dict()를 사용하여 체크포인트에서 재개할 수 있습니다:

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> for idx, example in enumerate(iterable_dataset):
...     print(example)
...     if idx == 2:
...         state_dict = iterable_dataset.state_dict()
...         print("checkpoint")
...         break
>>> iterable_dataset.load_state_dict(state_dict)
>>> print(f"restart from checkpoint")
>>> for example in iterable_dataset:
...     print(example)

반환 값:

{'a': 0}
{'a': 1}
{'a': 2}
checkpoint
restart from checkpoint
{'a': 3}
{'a': 4}
{'a': 5}

내부적으로, iterable dataset은 현재 읽는 중인 shard와 해당 shard 내의 예시 인덱스를 추적하며, 이 정보를 state_dict에 저장합니다.

체크포인트에서 재개하려면, 데이터셋은 이전에 읽은 모든 샤드를 건너뛰고 현재 샤드에서 다시 시작합니다. 그 다음 샤드를 읽고, 체크포인트의 정확한 예제에 도달할 때까지 예제를 건너뜁니다.

따라서 데이터셋을 다시 시작하는 것은 이미 반복된 샤드를 다시 읽지 않으므로 상당히 빠릅니다. 여전히 데이터셋을 재개하는 것은 일반적으로 즉각적이지 않으며, 현재 샤드의 처음부터 다시 읽기 시작하여 체크포인트 위치에 도달할 때까지 예제를 건너뛰어야 하기 때문입니다.

torchdata의 StatefulDataLoader와 함께 사용할 수 있습니다:

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(iterable_dataset, batch_size=32, num_workers=4)
>>> # checkpoint
>>> state_dict = dataloader.state_dict()  # uses iterable_dataset.state_dict() under the hood
>>> # resume from checkpoint
>>> dataloader.load_state_dict(state_dict)  # uses iterable_dataset.load_state_dict() under the hood

.shuffle()를 사용하는 경우를 제외하고, Resuming은 체크포인트가 저장된 위치에서 정확히 이어서 실행됩니다. Resuming 시 셔플 버퍼의 예제는 손실되며, 버퍼에는 새로운 데이터로 다시 채워집니다.

저장

반복 가능 데이터셋이 준비되면 Parquet 형식의 Hugging Face Dataset으로 저장하여 나중에 load_dataset()로 다시 사용할 수 있습니다.

Hugging Face의 데이터셋 저장소 이름을 push_to_hub()의 인수로 전달하여 데이터셋을 저장합니다. 이 방법은 데이터셋을 순회하면서 데이터를 Hugging Face에 점진적으로 업로드합니다:

dataset.push_to_hub("username/my_dataset")

기본적으로 IterableDataset.push_to_hub()은 dataset.num_shards개의 파일을 작성합니다. num_shards를 설정하여 다른 개수를 선택할 수 있습니다. 출력 샤드가 dataset.num_shards보다 많으면, 각 입력 샤드는 라운드 로빈 방식으로 예시를 출력 샤드 그룹에 작성하여 데이터에 대해 단일 패스로 분산됩니다. 요청된 샤드보다 예시가 적으면, 나머지 샤드는 빈 Parquet 파일이 됩니다.

데이터셋이 여러 개의 샤드로 구성되어 있는 경우(dataset.num_shards > 1)에는 여러 프로세스를 사용하여 병렬로 업로드할 수 있습니다. map() 또는 filter() 단계를 적용한 경우 특히 유용하며, 병렬 실행 시 더 빠르게 처리됩니다:

dataset.push_to_hub("username/my_dataset", num_proc=8)

load_dataset() 함수를 사용하여 데이터셋을 다시 로드합니다:

from datasets import load_dataset
reloaded_dataset = load_dataset("username/my_dataset")

내보내기

🤗 Datasets는 내보내기를 지원하므로 다른 애플리케이션에서 데이터셋을 사용할 수 있습니다. 다음 표는 현재 내보낼 수 있는 지원 파일 형식을 보여줍니다:

파일 형식Export method
CSVIterableDataset.to_csv()
JSONIterableDataset.to_json()
ParquetIterableDataset.to_parquet()
SQLIterableDataset.to_sql()
메모리 내 Python 객체IterableDataset.to_pandas(), IterableDataset.to_polars() or IterableDataset.to_dict()

예를 들어, 데이터셋을 다음과 같이 CSV 파일로 내보낼 수 있습니다:

>>> dataset.to_csv("path/of/my/dataset.csv")

대용량 데이터셋을 사용하는 경우, 샤드별로 파일을 저장할 수 있습니다. 예를 들어

>>> num_shards = dataset.num_shards
>>> for index in range(num_shards):
...     shard = dataset.shard(index, num_shards)
...     shard.to_parquet(f"path/of/my/dataset/data-{index:05d}.parquet")
Update on GitHub