Перевод с: Big data? 🤗 Datasets to the rescue!

В настоящее время нередко приходится работать с большими наборами данных ГБ, особенно для предварительного обучения моделей Transformer, таких как BERT или GPT-2, с нуля. В таких случаях даже загрузка данных может стать проблемой. Например, корпус WebText, используемый для предварительного обучения GPT-2, содержит более 8 миллионов документов и 40 ГБ текста — загрузка его в оперативную память вашего компьютера, вероятно, его испортит. К счастью, конструкция набора данных HuggingFace DataSet позволила преодолеть эти ограничения. Она решает проблему управления памятью, рассматривая набор данных как файл с отображением в памяти и обрабатывая записи в корпусе посредством потоковой передачи. Мы будем использовать огромный корпус размером 825 ГБ (называемый Pile) для изучения этих свойств набора данных HuggingFace.

Что такое Пайл?

Pile — это корпус английских текстов, созданный EleutherAI для обучения крупномасштабных языковых моделей. Он включает в себя различные наборы данных — от научных статей до репозиториев кода GitHub и отфильтрованного веб-текста. Учебный корпус предоставляется частями по 14 ГБ, также вы можете загрузить несколько отдельных компонентов. Начните с набора данных PubMed Abstracts — корпуса рефератов из 15 миллионов биомедицинских публикаций на PubMed. Набор данных имеет формат строк JSON и сжат с помощью библиотеки zstandard, поэтому сначала нам нужно его установить:

!pip install zstandard

Далее вы можете использовать метод загрузки набора данных, предоставленный HuggingFace, для загрузки:

from datasets import load_dataset

# This takes a few minutes to run, so go grab a tea or coffee while you wait :)
data_files = "https://the-eye.eu/public/AI/pile_preliminary_components/PUBMED_title_abstracts_2019_baseline.jsonl.zst"
pubmed_dataset = load_dataset("json", data_files=data_files, split="train")
pubmed_dataset

Dataset({
  features: ['meta', 'text'],
  num_rows: 15518009
})

Как видите, имеется 15518009 строк и 2 столбца данных. Вы можете просмотреть первый пример содержимого набора выходных данных ниже:

print pubmed_dataset[0];

# output:
{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

Это похоже на выдержку из медицинской статьи. Теперь давайте посмотрим, сколько оперативной памяти мы использовали для загрузки набора данных!

memory mapping

Простой способ измерить использование памяти в Python — использовать библиотеку psutil, которую можно установить с помощью pip следующим образом:

!pip install psutil

Он предоставляет класс Process, который позволяет нам проверять использование памяти текущим процессом следующим образом:

import psutil

# Process.memory_info is expressed in bytes, so convert to megabytes
print(f"RAM used: {psutil.Process().memory_info().rss / (1024 * 1024):.2f} MB")

# output:
RAM used: 5678.33 MB

Атрибут rss здесь относится к размеру резидентного набора, который представляет собой часть памяти, занимаемую процессом в оперативной памяти. Это измерение также включает в себя память, используемую интерпретатором Python и загруженными нами библиотеками, поэтому фактический объем памяти, используемый для загрузки набора данных, меньше. Для сравнения мы используем свойство dataset_size, чтобы увидеть размер набора данных на диске. Поскольку результат по-прежнему выражается в байтах, нам нужно вручную преобразовать его в гигабайты:

print(f"Number of files in dataset : {pubmed_dataset.dataset_size}")
size_gb = pubmed_dataset.dataset_size / (1024**3)
print(f"Dataset size (cache file) : {size_gb:.2f} GB")

# output
Number of files in dataset : 20979437051
Dataset size (cache file) : 19.54 GB

Очень хорошо, хоть и близко 20 GB большой, но мы можем использовать меньше RAM Загружайте и получайте доступ к наборам данных! если вы знакомы с Панды, этот результат может вас удивить, потому что Wes Kinney Знаменитое эмпирическое правило заключается в том, что обычно вам нужно RAM размер набора данных 5 приезжать 10 раз. Так как же набор данных HuggingFace решает эту проблему управления памятью? HuggingFace Datasets Рассматривая каждый набор данных как файл с отображением в памяти, он обеспечивает RAM Файловая система хранит сопоставления между элементами, позволяя библиотекам получать доступ к набору данных и выполнять операции с элементами без необходимости полной загрузки их в приезжать Память. Файлы, отображаемые в памяти, также могут совместно использоваться несколькими процессами, что делает Dataset.map() Подобные методы можно распараллелить без перемещения или копирования набора данных. Под капотом эти функции обеспечивает Apache Arrow формат памяти и pyarrow библиотека, которая делает загрузку и обработку данных молниеносной. Чтобы увидеть это в действии, давайте пройдемся по PubMed Abstracts Все элементы в наборе данных для запуска некоторых тестов скорости:

import timeit

code_snippet = """batch_size = 1000

for idx in range(0, len(pubmed_dataset), batch_size):
    _ = pubmed_dataset[idx:idx + batch_size]
"""

time = timeit.timeit(stmt=code_snippet, number=1, globals=globals())
print(
    f"Iterated over {len(pubmed_dataset)} examples (about {size_gb:.1f} GB) in "
    f"{time:.1f}s, i.e. {size_gb/time:.3f} GB/s"
)

# output:
'Iterated over 15518009 examples (about 19.5 GB) in 64.2s, i.e. 0.304 GB/s'

Здесь мы используем модуль Python timeit для измерения времени выполнения code_snippet. Обычно вы можете получить это за несколько десятых ГБ/секунду Сколько приезжать Перебор наборов данных со скоростью ГБ/сек. Это отлично работает для подавляющего большинства приложений, но иногда вам приходится работать с наборами данных, которые слишком велики для хранения на жестком диске вашего ноутбука. Например, если мы попытаемся загрузить весь Куча, нам понадобится 825 GB доступного дискового пространства! Чтобы справиться с такими ситуациями, обнимайте Face Datasets Предоставляется функция потоковой передачи, позволяющая нам динамически загружать элементы и получать к ним доступ, не загружая весь набор данных.

Streaming Datasets

Включение потоковой передачи набора данных,Вам просто нужноStreaming=TrueПараметры, переданные вload_dataset()функция。 Например, давайте снова загрузим PubMed Abstracts Набор данных, но в потоковом режиме:

pubmed_dataset_streamed = load_dataset(
    "json", data_files=data_files, split="train", streaming=True
)

Streaming=True Возвращаемый из объекта нам незнакомый вид, встречающийся в других разделах этой главы. Набор данных, но IterableDataset。 Как следует из названия, посетить IterableDataset элементы, нам нужно перебрать их. Мы можем получить доступ к первому элементу набора потоковых данных следующим образом:

next(iter(pubmed_dataset_streamed))

# output
{'meta': {'pmid': 11409574, 'language': 'eng'},
 'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection.\nTo determine the prevalence of hypoxaemia in children aged under 5 years suffering acute lower respiratory infections (ALRI), the risk factors for hypoxaemia in children under 5 years of age with ALRI, and the association of hypoxaemia with an increased risk of dying in children of the same age ...'}

Можно использоватьIterableDataset.map()Мгновенный поток обработкиданныеконцентрированныйизэлемент,Если вам нужно пометить ввод,Это существование очень полезно во время тренировок.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
tokenized_dataset = pubmed_dataset_streamed.map(lambda x: tokenizer(x["text"]))
next(iter(tokenized_dataset))

# output
{'input_ids': [101, 4958, 5178, 4328, 6779, ...], 'attention_mask': [1, 1, 1, 1, 1, ...]}

Также доступенIterableDataset.shuffle()Тип конвекцииданные Наборы перемешиваются,Но с Dataset.shuffle() Другое дело, он работает только с предопределенными buffer_size Перетасуйте элементы:

shuffled_dataset = pubmed_dataset_streamed.shuffle(buffer_size=10_000, seed=42)
next(iter(shuffled_dataset))

# output
{'meta': {'pmid': 11410799, 'language': 'eng'},
 'text': 'Randomized study of dose or schedule modification of granulocyte colony-stimulating factor in platinum-based chemotherapy for elderly patients with lung cancer ...'}

В этом примере мы начинаем с предыдущего 10,000 Из примеров случайным образом выбирается один пример. После доступа к примеру его позиция в буфере заполняется следующим примером в корпусе (т. е. первым примером в приведенном выше случае). 10,001 примеры). Также доступен IterableDataset.take() иIterableDataset.skip()функцияиз потокового вещанияданныеконцентрированныйвыбиратьэлемент,Его принцип действия такой же, как и уDataset.select()похожий。 Например, чтобы выбрать PubMed Abstracts перед набором данных 5 В качестве примера мы можем сделать следующее:

dataset_head = pubmed_dataset_streamed.take(5)
list(dataset_head)

# output
[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'pmid': 11409575, 'language': 'eng'},
  'text': 'Clinical signs of hypoxaemia in children with acute lower respiratory infection: indicators of oxygen therapy ...'},
 {'meta': {'pmid': 11409576, 'language': 'eng'},
  'text': "Hypoxaemia in children with severe pneumonia in Papua New Guinea ..."},
 {'meta': {'pmid': 11409577, 'language': 'eng'},
  'text': 'Oxygen concentrators and cylinders ...'},
 {'meta': {'pmid': 11409578, 'language': 'eng'},
  'text': 'Oxygen supply in rural africa: a personal experience ...'}]

Аналогично, разделения обучения и проверки можно создать из перетасованного набора данных с помощью функции IterableDataset.skip() следующим образом:

# Skip the first 1,000 examples and include the rest in the training set
train_dataset = shuffled_dataset.skip(1000)
# Take the first 1,000 examples for the validation set
validation_dataset = shuffled_dataset.take(1000)

Давайте завершим наше исследование потоков наборов данных с помощью общего приложения: объединения нескольких наборов данных для создания единого корпуса. 🤗 Datasets предоставил interleave_datasets() функция, она будет IterableDataset Преобразовать список объектов в один IterableDataset, где элементы нового набора данных получаются путем чередования исходных примеров. Эта функция особенно полезна при попытке объединить большие наборы данных, поэтому мы выполняем потоковую передачу. Pile из FreeLaw Например, подмножество из судов США из 51 GB Набор данных юридического заключения:

law_dataset_streamed = load_dataset(
    "json",
    data_files="https://the-eye.eu/public/AI/pile_preliminary_components/FreeLaw_Opinions.jsonl.zst",
    split="train",
    streaming=True,
)
next(iter(law_dataset_streamed))

{'meta': {'case_ID': '110921.json',
  'case_jurisdiction': 'scotus.tar.gz',
  'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}

Этот набор данных достаточно велик для большинства ноутбуков. RAM Это стресс, но мы смогли загрузить его и получить к нему доступ без каких-либо проблем! Теперь давайте FreeLaw и PubMed Abstracts концентрация данных из примера с interleave_datasets() Функции в сочетании:

from itertools import islice
from datasets import interleave_datasets

combined_dataset = interleave_datasets([pubmed_dataset_streamed, law_dataset_streamed])
list(islice(combined_dataset, 2))

[{'meta': {'pmid': 11409574, 'language': 'eng'},
  'text': 'Epidemiology of hypoxaemia in children with acute lower respiratory infection ...'},
 {'meta': {'case_ID': '110921.json',
   'case_jurisdiction': 'scotus.tar.gz',
   'date_created': '2010-04-28T17:12:49Z'},
  'text': '\n461 U.S. 238 (1983)\nOLIM ET AL.\nv.\nWAKINEKONA\nNo. 81-1581.\nSupreme Court of United States.\nArgued January 19, 1983.\nDecided April 26, 1983.\nCERTIORARI TO THE UNITED STATES COURT OF APPEALS FOR THE NINTH CIRCUIT\n*239 Michael A. Lilly, First Deputy Attorney General of Hawaii, argued the cause for petitioners. With him on the brief was James H. Dannenberg, Deputy Attorney General...'}]

Здесь мы используем Python из itertools В модуле islice() Функция Выбрав первые два примера из объединенного набора данных, мы можем увидеть, что они соответствуют первому примеру из двух исходных наборов данных.

Наконец, если вы хотите транслировать все 825 ГБ из Pile, все готовые файлы можно получить следующим образом:

base_url = "https://the-eye.eu/public/AI/pile/"
data_files = {
    "train": [base_url + "train/" + f"{idx:02d}.jsonl.zst" for idx in range(30)],
    "validation": base_url + "val.jsonl.zst",
"test": base_url + "test.jsonl.zst",
}
pile_dataset = load_dataset("json", data_files=data_files, streaming=True)
next(iter(pile_dataset["train"]))

# output
{'meta': {'pile_set_name': 'Pile-CC'},
 'text': 'It is done, and submitted. You can play “Survival of the Tastiest” on Android, and on the web...'}

Подвести итог

Подвести итог Приходите и посмотрите,В основном за счет картографирования Памяти и потоковой обработки для достижения загрузки большого набора данных.,Это также широко используемое решение в отрасли.