Работа с датасетами: load_dataset и map

Учимся загружать и подготавливать данные для моделей через библиотеку datasets.

datasets — библиотека Hugging Face для загрузки, хранения и обработки наборов данных, оптимизированная под большие объёмы.

Загрузка одной строкой

Любой датасет с Hub загружается функцией load_dataset по имени. Он скачивается, кэшируется и разбивается на части (train/test/validation).

from datasets import load_dataset

ds = load_dataset("imdb")
print(ds)
# DatasetDict({
#   train: Dataset(features: ['text','label'], num_rows: 25000),
#   test:  Dataset(features: ['text','label'], num_rows: 25000)
# })
print(ds["train"][0]["label"])  # 0 или 1

Преобразование через map

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

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True)

tokenized = ds.map(tokenize, batched=True)

Имитация map на чистом Python

Чтобы прочувствовать идею map, посмотрим на аналог из стандартного Python (реально запускается):

data = [{"text": "hello"}, {"text": "hi there"}]

def add_length(example):
    example["len"] = len(example["text"].split())
    return example

processed = [add_length(x) for x in data]
for row in processed:
    print(row)

Вывод:

{'text': 'hello', 'len': 1}
{'text': 'hi there', 'len': 2}

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

Как работает под капотом

datasets хранит данные в формате Apache Arrow — колоночном бинарном формате, который позволяет работать с датасетами больше оперативной памяти (memory mapping). Результаты map кэшируются на диск по хешу функции и аргументов, поэтому повторный запуск того же преобразования мгновенный. Параметр batched=True передаёт в функцию сразу группу примеров, что заметно быстрее поэлементной обработки.

Частые ошибки

Грузить весь датасет в обычный список. Теряется ленивость и кэширование Arrow; на больших данных кончится память.
Забыть truncation при токенизации. Длинные тексты превысят лимит модели и вызовут ошибку.
Менять функцию map, не очистив кэш. Иногда нужен load_from_cache_file=False.

Итог

load_dataset("имя") скачивает и кэширует датасет с разбивкой на части.
map применяет функцию ко всем примерам, обычно для токенизации.
batched=True ускоряет обработку.
Под капотом — формат Arrow и кэширование на диск.