Препарируем 🤗 Datasets

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

Управление данными

Как и в Pandas, 🤗 Datasets предоставляет несколько функция для управления содержимым объектов Dataset и DatasetDict. Мы уже познакомились с методом Dataset.map() в главе 3, а далее мы посмотрим на другие функции, имеющиеся в нашем распоряжении.

Для этого примера мы будем использовать датасет Drug Review Dataset, расположенный на сервере UC Irvine Machine Learning Repository и содержащий отзывы пациентов на различные лекарства, сведения о состоянии пациентов и рейтинг удовлетворенности, выраженный в 10-балльной шкале.

Для начала необходимо скачать и разархивировать датасет, мы используем для этого команды wget и unzip:

!wget "https://archive.ics.uci.edu/ml/machine-learning-databases/00462/drugsCom_raw.zip"
!unzip drugsCom_raw.zip

Файл TSV - это просто разновидность CSV файла, содержащий табуляции вместо запятых в качестве разделителя, а значит мы можем его загрузить с помощью скрипта csv и аргумента delimiter через функцию load_dataset():

from datasets import load_dataset

data_files = {"train": "drugsComTrain_raw.tsv", "test": "drugsComTest_raw.tsv"}
# \t is the tab character in Python
drug_dataset = load_dataset("csv", data_files=data_files, delimiter="\t")

Хорошей практикой при исследовании данных является взятие небольшого случайного подмножества для понимания типов данных и их особенностей. В библиотеке 🤗 Datasets мы можем сделать случайную выборку путем последовательного вызова функций Dataset.shuffle() и Dataset.select():

drug_sample = drug_dataset["train"].shuffle(seed=42).select(range(1000))
# Peek at the first few examples
drug_sample[:3]

{'Unnamed: 0': [87571, 178045, 80482],
 'drugName': ['Naproxen', 'Duloxetine', 'Mobic'],
 'condition': ['Gout, Acute', 'ibromyalgia', 'Inflammatory Conditions'],
 'review': ['"like the previous person mention, I'm a strong believer of aleve, it works faster for my gout than the prescription meds I take. No more going to the doctor for refills.....Aleve works!"',
  '"I have taken Cymbalta for about a year and a half for fibromyalgia pain. It is great\r\nas a pain reducer and an anti-depressant, however, the side effects outweighed \r\nany benefit I got from it. I had trouble with restlessness, being tired constantly,\r\ndizziness, dry mouth, numbness and tingling in my feet, and horrible sweating. I am\r\nbeing weaned off of it now. Went from 60 mg to 30mg and now to 15 mg. I will be\r\noff completely in about a week. The fibro pain is coming back, but I would rather deal with it than the side effects."',
  '"I have been taking Mobic for over a year with no side effects other than an elevated blood pressure.  I had severe knee and ankle pain which completely went away after taking Mobic.  I attempted to stop the medication however pain returned after a few days."'],
 'rating': [9.0, 3.0, 10.0],
 'date': ['September 2, 2015', 'November 7, 2011', 'June 5, 2013'],
 'usefulCount': [36, 13, 128]}

Заметьте, что мы зафикисировали переменную seed для воспроизводимости результатов. Dataset.select() ожидает на вход итерируемый объект, содержащий индексы, поэтому мы передали range(1000) для взятия первых 1000 объектов перемешанного датасета. Для этой подвыборки мы можем сразу увидеть некоторые особенности в данных:

• Колонка Unnamed: 0 выглядит как обезличенный ID для каждого пациента.
• Колонка condition включает в себя смесь лейблов в нижнем и верхнем регистре.
• Отзывы переменной длины и содержат смесь разделителей текста (\r\n) и HTML-кодов (например, &\#039;).

Давайте посмотрим, как мы можем использовать 🤗 Datasets для обработки этих особенностей. Чтобы проверить, что наша гипотеза об уникальности справедлива, мы можем использовать функцию Dataset.unique() для проверки, что число ID совпадает с числом строк в обоих датасетах (обучающем и тестовом):

for split in drug_dataset.keys():
    assert len(drug_dataset[split]) == len(drug_dataset[split].unique("Unnamed: 0"))

По всей видимости, наша гипотеза подтвердилась, так что перейдем к очистке датасета. Для начала переименуем Unnamed: 0 во что-то более интерпретируемое. Мы можем использовать функцию DatasetDict.rename_column() для переименования столбцы на обоих сплитах (обучающем и тестовом):

drug_dataset = drug_dataset.rename_column(
    original_column_name="Unnamed: 0", new_column_name="patient_id"
)
drug_dataset

DatasetDict({
    train: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount'],
        num_rows: 161297
    })
    test: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount'],
        num_rows: 53766
    })
})

Далее нормализуем все лейблы столбца condition с применением Dataset.map(). Так же, как мы делали токенизацию в главе 3, мы можем определить простую функцию, которая будет применения для всех строк каждого сплита в drug_dataset:

def lowercase_condition(example):
    return {"condition": example["condition"].lower()}


drug_dataset.map(lowercase_condition)

AttributeError: 'NoneType' object has no attribute 'lower'

О нет! При запуске этой функции мы столкнулись с проблемой! Из ошибки мы можем сделать вывод, что некоторые записи в колонке condition являются None, которые не могут быть приведены к нижнему регистру как обычные строковые типы данных. Давайте удалим эти строки с помощью Dataset.filter(), которая работает схожим с Dataset.map() образом и принимает на вход один экземпляр датасета. Вместо реализации собственной функции:

def filter_nones(x):
    return x["condition"] is not None

и вызова этой функции drug_dataset.filter(filter_nones), мы можем сделать то же самое с помощью lambda-функции. В Python лямбда-функции - это небольшие функции, которые вы можете определить без явного их именования. Общий вид, которым их можно задать:

lambda <arguments> : <expression>

где lambda - одно из ключевых слов Python, а <arguments> - список или множество разделенных запятой значений, которые пойдут на вход функции, и <expression> задает операции, которые вы хотите применить к аргументам. Например, мы можем задать простую лямбда-функцию, которая возводит в квадрат числа:

lambda x : x * x

Чтобы применить эту функцию, мы должны заключить ее и аргументы в скобки:

(lambda x: x * x)(3)

9

По аналогии мы можем задать лямбда-функцию с несколькими аргументами, которые необходимо разделить запятыми. Например, мы можем вычислить площадь треугольника следующим образом:

(lambda base, height: 0.5 * base * height)(4, 8)

16.0

Лямбда-функции удобны, когда вы хотите определить маленькие одноразовые функции (для более подробной информации об этих функциях мы рекомендуем изучить превосходную публикацию Real Python tutorial за авторством Andre Burgaud). В контексте библиотеки 🤗 Datasets мы можем использовать лямбда-функции для задания простых операций map и filter, давайте попробуем устранить None-записи из нашего датасета:

drug_dataset = drug_dataset.filter(lambda x: x["condition"] is not None)

После удаления None записей, мы можем нормализовать колонку condition:

drug_dataset = drug_dataset.map(lowercase_condition)
# Check that lowercasing worked
drug_dataset["train"]["condition"][:3]

['left ventricular dysfunction', 'adhd', 'birth control']

Заработало! Сейчас мы очистили лейблы, давайте теперь посмотрим на то, как можно очистить непосредственно отзывы.

Создание новых столбцов

Всякий раз, когда вы имеете дело с отзывами клиентов, хорошей практикой является проверка количества слов в каждом отзыве. Обзор может состоять всего из одного слова, например «Отлично!» или быть полномасштабным эссе с тысячами слов, и в зависимости от варианта использования вам нужно будет по-разному справляться с этими случаями. Чтобы вычислить количество слов в каждом обзоре, мы будем использовать грубую эвристику, основанную на разбиении каждого текста по пробелам.

Зададим простую функцию, которая вычисляет число слов в каждом отзыве:

def compute_review_length(example):
    return {"review_length": len(example["review"].split())}

В отличие от функции lowercase_condition(), compute_review_length() возвращает словарь, чьи ключи не соответствуют ни одному названию колонки в нашем датасете. В этом случае при исполнении compute_review_length() (переданного в Dataset.map()) функция будет применена ко всем строкам в датасете и создаст новый столбец с именем review_length:

drug_dataset = drug_dataset.map(compute_review_length)
# Посмотрим на первый объект обучающей части датасета
drug_dataset["train"][0]

{'patient_id': 206461,
 'drugName': 'Valsartan',
 'condition': 'left ventricular dysfunction',
 'review': '"It has no side effect, I take it in combination of Bystolic 5 Mg and Fish Oil"',
 'rating': 9.0,
 'date': 'May 20, 2012',
 'usefulCount': 27,
 'review_length': 17}

Как и ожадалось, мы видим колонку с именем review_length, которая добавлена к нашему обучающему датасету. Мы можем отсортировать по этой колонке наш датасет с помощью функции Dataset.sort() и посмотреть на «экстремальные» значения:

drug_dataset["train"].sort("review_length")[:3]

{'patient_id': [103488, 23627, 20558],
 'drugName': ['Loestrin 21 1 / 20', 'Chlorzoxazone', 'Nucynta'],
 'condition': ['birth control', 'muscle spasm', 'pain'],
 'review': ['"Excellent."', '"useless"', '"ok"'],
 'rating': [10.0, 1.0, 6.0],
 'date': ['November 4, 2008', 'March 24, 2017', 'August 20, 2016'],
 'usefulCount': [5, 2, 10],
 'review_length': [1, 1, 1]}

Как и ожидалось, некоторые отзывы содержат одно слово, хотя это и может быть допустимо для задачи оценки тональности текста, вряд ли будет полезно если мы хотим предсказывать состояние пациента.

Давайте применим функцию Dataset.filter() для удаления отзывов, содержащих меньше 30 слов. Схожим образом мы применяли её для столбца condition: мы можем отфильтровать отзывы, в которых число слов меньше порога:

drug_dataset = drug_dataset.filter(lambda x: x["review_length"] > 30)
print(drug_dataset.num_rows)

{'train': 138514, 'test': 46108}

Как вы можете увидеть, эта функция удалила около 15% отзывов из наших исходных обучающих и тестовых наборов данных.

Последняя вещь, которую нам необходимо сделать, это справиться с присутствием HTML-кодами символов в наших отзывах. Мы можем использовать модуль html и метод unescape() чтобы избавиться от них:

import html

text = "I'm a transformer called BERT"
html.unescape(text)

"I'm a transformer called BERT"

Для этого будем использовать Dataset.map() на всем нашем корпусе текстов:

drug_dataset = drug_dataset.map(lambda x: {"review": html.unescape(x["review"])})

Как видите, метод Dataset.map() крайне полезен для препроцессинга данных — хотя мы и воспользовались только малой частью его возможностей!

Суперспособности метода map()

Метод Dataset.map() принимает аргумент batched, который, если установлен в значение True, заставляет его сразу отправлять батч элементов в функцию map() (размер батча можно настроить, но по умолчанию он равен 1000). Например, предыдущая функция map(), которая экранировала весь HTML-код, требовала некоторого времени для запуска (вы можете узнать время взглянув на индикаторы выполнения процесса). Мы можем ускорить это, обрабатывая несколько элементов одновременно, используя list comprehension.

Когда вы указываете batched=True, функция получает словарь с полями набора данных, но каждое значение теперь представляет собой список значений, а не просто одно значение. Возвращаемое значение Dataset.map() должно быть одинаковым: словарь с полями, которые мы хотим обновить или добавить в наш набор данных, и список значений. Например, вот еще один способ устранить все символы HTML, но с использованием batched=True:

new_drug_dataset = drug_dataset.map(
    lambda x: {"review": [html.unescape(o) for o in x["review"]]}, batched=True
)

Если вы запустите этот код в блокноте, вы увидите, что эта команда выполняется намного быстрее, чем предыдущая. И это не потому, что наши отзывы уже были HTML-экранированными — если вы повторно выполните инструкцию из предыдущего раздела (без batched=True), это займет столько же времени, сколько и раньше. Это связано с тем, что обработка списков обычно выполняется быстрее, чем выполнение того же кода в цикле for, мы также повышаем производительность за счет одновременного доступа к множеству элементов, а не по одному.

Использование Dataset.map() с batched=True – хороший способ «разблокировать» скоростные ограничения “быстрых” токенизаторов, с которыми мы познакомимся в главе 6, которые могут быстро токенизировать большие списки текста. Например, чтобы токенизировать все отзывы на лекарства с помощью быстрого токенизатора, мы можем использовать функцию, подобную этой:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")


def tokenize_function(examples):
    return tokenizer(examples["review"], truncation=True)

Как вы видели в главе 3, мы можем передать один или несколько элементов в токенизатор, так что мы можем использовать эту функцию без параметра batched=True. Давайте воспользуемся этой возможностью и сравним производительность. В ноутбуке можно замерить время выполнения функции путем добавления %time перед строкой кода, время исполнения которой вы хотите измерить:

%time tokenized_dataset = drug_dataset.map(tokenize_function, batched=True)

Также присутствует возможность измерить время выполнения всей ячейки: нужно заменить %time на %%time в начале ячейки. На нашем оборудовании это заняло 10.8 секунд. Это значение расположено после слов “Wall time”.

Вот результаты, которые мы получили без и с применением батчинга, и двумя разными по скорости токенизаторами:

По результатам видно, что использование быстрого токенизатора с параметром batched=True приводит к ускорению выполнения в 30 раз – это потрясающе! Это главная причина, почему быстрые токенизаторы применяются по умолчанию при использовании класса AutoTokenizer (и почему они называются “быстрыми”). Возможность достичь такой скорости выполнения достигается засчет исполнения кода токенизаторов на языке Rust, который легко позволяет распараллелить выполнение кода.

Параллелизация также позволяет почти в 6 раз ускорить быстрые токенизаторы с использованием batched=True: вы не можете пареллелизовать едничную операцию токенизации, но когда вы токенизируете много различных текстов одновременно, вы можете распределить выполнение на несколько процессов, каждый из которых будет отвечать за собственный текст.

Dataset.map() также обладает возможностями параллелизации. Поскольку метод не реализован на Rust, он не позволят “медленному” токенизатору “догнать” быстрый, но все же может быть полезен (особенно если вы используете токенизатор, у которого нет быстрой версии). Чтобы включить многопроцессорность, используйте аргумент num_proc и укажите количество процессов, которые будут использоваться в вашем вызове Dataset.map():

slow_tokenizer = AutoTokenizer.from_pretrained("bert-base-cased", use_fast=False)


def slow_tokenize_function(examples):
    return slow_tokenizer(examples["review"], truncation=True)


tokenized_dataset = drug_dataset.map(slow_tokenize_function, batched=True, num_proc=8)

Вы можете поэкспериментировать и выяснить, какое число num_proc даст наилучший результат, в нашем случае значение 8 стало оптимальным. Вот значения, которые мы получили с и без использования мультипроцессинга:

Это гораздо более разумные результаты для “медленного” токенизатора, но производительность быстрого токенизатора также существенно выросла. Однако, обратите внимание, что это не всегда так — для значений num_proc, отличных от 8, наши тесты показали, что быстрее использовать batched=True без этой опции. Как правило, мы не рекомендуем использовать мультипроцессинг Python для “быстрых” токенизаторов с параметром batched=True.

Объединение всей этой функциональности во всего лишь один метод само по себе прекрасно, но это еще не все! Используя Dataset.map() и batched=True вы можете поменять число элементов в датасете. Это очень полезно во множестве ситуаций, например, когда вы хотите создать несколько обучающих признаков из одного экземпляра текста. Мы воспользуеся этой возможностью на этапе препроцессинга для нескольких NLP-задач, которые рассмотрим в главе 7

Давайте посмотрим как это работает! В этом примере мы токенизируем наши тексты и обрежем их до максимальной длины в 128, однако мы попросим токенизатор вернуть нам все получившиеся токены, а не только начальные. Это может быть сделано с помощью параметра return_overflowing_tokens=True:

def tokenize_and_split(examples):
    return tokenizer(
        examples["review"],
        truncation=True,
        max_length=128,
        return_overflowing_tokens=True,
    )

Давайте протестируем это на одном тексте прежде, чем использовать Dataset.map() на всем датасете:

result = tokenize_and_split(drug_dataset["train"][0])
[len(inp) for inp in result["input_ids"]]

[128, 49]

Итак, наш первый текст в обучающей части выборки стал состоять из двух признаков, т.к. токенизатор токенизировал не только первые 128 элементов, но и оставшиеся 49 тоже. Давайте применим токенизатор ко всем элементам датасета!

tokenized_dataset = drug_dataset.map(tokenize_and_split, batched=True)

ArrowInvalid: Column 1 named condition expected length 1463 but got length 1000

О, нет! Не сработало! Почему? Посмотрим на ошибку: несовпадение в длинах, один из которых длиной 1463, а другой – 1000. Если вы обратитесь в документацию Dataset.map(), вы можете увидеть, что одно из этих чисел – число объектов, поданных на вход функции, а другое –

Проблема заключается в том, что мы пытаемся смешать два разных датасета разной размерности: число колонок датасета drug_dataset равняется 1000, а нужный нам tokenized_dataset имеет 1463 колонки. Чтобы избежать этой ошибки, необходимо удалить несколько столбцов из старого датасета и сделать оба датасета одинакового размера. Мы можем достичь этого с помощью аргумента remove_columns:

tokenized_dataset = drug_dataset.map(
    tokenize_and_split, batched=True, remove_columns=drug_dataset["train"].column_names
)

Теперь это работает без ошибок. Мы можем проверить, что наш новый датасет имеет бОльшее число элементов, просто сравним длины датасетов:

len(tokenized_dataset["train"]), len(drug_dataset["train"])

(206772, 138514)

Мы упоминали о том, что мы можем справиться с несовпадением длин путем исправления числа колонок старого датасета. Чтобы сделать это, нам необходимо получить поле overflow_to_sample_mapping, которое вернет нам токенизатор, если мы зададим аргумент return_overflowing_tokens=True. Это даст нам необходимое соответствие между индексом новых и старых признаков. После этого мы сможем ассоциировать каждый ключ нашего оригинального датасета со списком значений нужного размера, повторяя значения каждого примера столько раз, сколько он генерирует новые функции:

Для этого нам понадобится поле overflow_to_sample_mapping, которое возвращает токенизатор, когда мы устанавливаем return_overflowing_tokens=True. Это дает нам сопоставление индекса новой функции с индексом выборки, из которой он произошел. Используя это, мы можем связать каждый ключ, присутствующий в нашем исходном наборе данных, со списком значений нужного размера, повторяя значения каждого примера столько раз, сколько он генерирует новые функции:

def tokenize_and_split(examples):
    result = tokenizer(
        examples["review"],
        truncation=True,
        max_length=128,
        return_overflowing_tokens=True,
    )
    # Extract mapping between new and old indices
    sample_map = result.pop("overflow_to_sample_mapping")
    for key, values in examples.items():
        result[key] = [values[i] for i in sample_map]
    return result

Мы можем убедиться, что это сработало в Dataset.map() и без удаления старых столбцов:

tokenized_dataset = drug_dataset.map(tokenize_and_split, batched=True)
tokenized_dataset

DatasetDict({
    train: Dataset({
        features: ['attention_mask', 'condition', 'date', 'drugName', 'input_ids', 'patient_id', 'rating', 'review', 'review_length', 'token_type_ids', 'usefulCount'],
        num_rows: 206772
    })
    test: Dataset({
        features: ['attention_mask', 'condition', 'date', 'drugName', 'input_ids', 'patient_id', 'rating', 'review', 'review_length', 'token_type_ids', 'usefulCount'],
        num_rows: 68876
    })
})

Мы получаем то же количество признаков, что и раньше, но здесь мы сохранили все старые поля. Если они вам нужны для некоторой постобработки после применения вашей модели, вы можете использовать этот подход.

Теперь вы видели, как 🤗 Datasets можно использовать для предварительной обработки набора данных различными способами. Хотя функции обработки 🤗 Datasets покроют большую часть ваших потребностей в обучении модели, могут быть случаи, когда вам нужно будет переключиться на Pandas, чтобы получить доступ к более мощным функциям, таким как DataFrame.groupby() или API высокого уровня для визуализации. К счастью, 🤗 Datasets предназначены для взаимодействия с такими библиотеками, как Pandas, NumPy, PyTorch, TensorFlow и JAX. Давайте посмотрим, как это работает.

От Dataset а к DataFrame ам и назад

Для включения конвертации между различными библиотеками 🤗 Datasets предоставляет функцию Dataset.set_format(). Эта функция только изменяет выходной формат датасета, так что вы можете переключиться на другой формат не изменяя саму структуру данных, которая остается Apache Arrow. Смена формата происходит in place. Для демонстрации давайте попробуем сконвертировать наш датасет в формат Pandas:

drug_dataset.set_format("pandas")

Теперь при обращении к элементам датасета мы будем получать pandas.DataFrame вместо словаря:

drug_dataset["train"][:3]

Давайте создадим pandas.DataFrame для всего обучающего множества, выбрав все элементы из drug_dataset["train"]:

train_df = drug_dataset["train"][:]

Начиная с этого момента мы можем использовать всю функциональность Pandas. Например, мы можем иначе посчитать расределение condition среди нашей выборки:

frequencies = (
    train_df["condition"]
    .value_counts()
    .to_frame()
    .reset_index()
    .rename(columns={"index": "condition", "condition": "frequency"})
)
frequencies.head()

И как только мы закончим наш анализ Pandas, мы всегда можем создать новый объект Dataset с помощью функции Dataset.from_pandas() следующим образом:

from datasets import Dataset

freq_dataset = Dataset.from_pandas(frequencies)
freq_dataset

Dataset({
    features: ['condition', 'frequency'],
    num_rows: 819
})

На этом мы заканчиваем наш обзор различных техник препроцессинга, доступных в 🤗 Datasets. Чтобы завершить этот раздел, давайте создадим валидационную часть выборки. Прежде, чем сделать это, мы сбросим формат drug_dataset обратно к "arrow":

drug_dataset.reset_format()

Создание валидационной выборки

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

🤗 Наборы данных предоставляют функцию Dataset.train_test_split(), основанную на известной функциональности из scikit-learn. Давайте используем её, чтобы разделить наш обучающий датасет непосредственно на обучающий и валидационный (мы устанавливаем аргумент seed для воспроизводимости):

drug_dataset_clean = drug_dataset["train"].train_test_split(train_size=0.8, seed=42)
# Переименуем "test"  в "validation"
drug_dataset_clean["validation"] = drug_dataset_clean.pop("test")
# Добавим "test" в наш `DatasetDict`
drug_dataset_clean["test"] = drug_dataset["test"]
drug_dataset_clean

DatasetDict({
    train: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length', 'review_clean'],
        num_rows: 110811
    })
    validation: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length', 'review_clean'],
        num_rows: 27703
    })
    test: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length', 'review_clean'],
        num_rows: 46108
    })
})

Отлично, теперь мы подготовили датасет, на котором можно обучить некоторые модели. В разделе 5 мы покажем, как загрузить датасеты на Hugging Face Hub, а пока закончим наш обзор и посмотрим несколько способов сохранения датасетов на локальный компьютер.

Сохранение датасетов

Несмотря на то, что 🤗 Datasets будут кэшировать все загруженные датасеты и операции, которые над ними выполняются, будут случаи, когда вам будет необходимо сохранить датасет на диск (например, если кэш был очищен). Как показано в таблице ниже, 🤗 Datasets предоставляет три главных функции для сохранения датасета в разных форматах.

Для примера сохраним наш очищенный датасет в формате Arrow:

drug_dataset_clean.save_to_disk("drug-reviews")

Эта функция создаст директорию следующей структуры:

drug-reviews/
├── dataset_dict.json
├── test
│   ├── dataset.arrow
│   ├── dataset_info.json
│   └── state.json
├── train
│   ├── dataset.arrow
│   ├── dataset_info.json
│   ├── indices.arrow
│   └── state.json
└── validation
    ├── dataset.arrow
    ├── dataset_info.json
    ├── indices.arrow
    └── state.json

где мы можем увидеть каждый сплит данных, ассоциированный с собственной таблицей dataset.arrow, и некоторыми метаданными, хранящимися в файлах dataset_info.json и state.json. Вы можете рассматривать формат Arrow просто как таблицу, которая оптимизирована для построения высокопроизводительных приложений для обработки и передачи больших датасетов.

После сохранения датасета мы можем загрузить его с использованием функции load_from_disk() следующим образом:

from datasets import load_from_disk

drug_dataset_reloaded = load_from_disk("drug-reviews")
drug_dataset_reloaded

DatasetDict({
    train: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length'],
        num_rows: 110811
    })
    validation: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length'],
        num_rows: 27703
    })
    test: Dataset({
        features: ['patient_id', 'drugName', 'condition', 'review', 'rating', 'date', 'usefulCount', 'review_length'],
        num_rows: 46108
    })
})

Для форматов CSV и JSON мы должны сохранять каждый сплит как отдельный файл. Один из способов это сделать – проитерироваться по ключам и значениям в объекте DatasetDict:

for split, dataset in drug_dataset_clean.items():
    dataset.to_json(f"drug-reviews-{split}.jsonl")

Этот код сохранит каждый блок нашего датасета в формате JSON Lines, где каждая строка будет сохранена как JSON-объект. Вот как будет выглядеть первый элемент нашей выборки:

!head -n 1 drug-reviews-train.jsonl

{"patient_id":141780,"drugName":"Escitalopram","condition":"depression","review":"\"I seemed to experience the regular side effects of LEXAPRO, insomnia, low sex drive, sleepiness during the day. I am taking it at night because my doctor said if it made me tired to take it at night. I assumed it would and started out taking it at night. Strange dreams, some pleasant. I was diagnosed with fibromyalgia. Seems to be helping with the pain. Have had anxiety and depression in my family, and have tried quite a few other medications that haven't worked. Only have been on it for two weeks but feel more positive in my mind, want to accomplish more in my life. Hopefully the side effects will dwindle away, worth it to stick with it from hearing others responses. Great medication.\"","rating":9.0,"date":"May 29, 2011","usefulCount":10,"review_length":125}

Мы можем использовать приёмы из раздела 2 для загрузки JSON-файлов:

data_files = {
    "train": "drug-reviews-train.jsonl",
    "validation": "drug-reviews-validation.jsonl",
    "test": "drug-reviews-test.jsonl",
}
drug_dataset_reloaded = load_dataset("json", data_files=data_files)

Вот и все, что нужно для нашего экскурса при работе с 🤗 Datasets! Мы очистили датасет для обучения модели, вот некоторые идеи, которые вы могли бы реализовать самостоятельно:

1. Примените знания из раздела 3 для обучения классификатора, который может предсказывать состояние пациента по отзыву на лекарство.
2. Используйте pipeline summarization из раздела 1для генерации саммари отзывов.

Далее мы посмотрим, как 🤗 Datasets могут помочь вам в работе с громадными датасетами, которые невозможно обработать на вашем ноутбуке!