Biblioteca Datasets de Hugging Face: Cargar y Procesar Datos

1. Descripción general

Cuando trabajes con modelos NLP o hagas fine-tuning de LLM, siempre necesitas un paso previo: preparar los datos en el formato correcto, asegurarte de que estén limpios y que se carguen correctamente en el pipeline de entrenamiento. La biblioteca datasets de Hugging Face resuelve exactamente este problema: te proporciona una API unificada para cargar datasets desde el Hub, procesarlos en lote mediante map/filter, y guardarlos en un formato optimizado para lecturas repetidas.

Después de leer este artículo, serás capaz de cargar datasets desde Hugging Face Hub, comprender la estructura de Dataset/DatasetDict, aprender a preprocesar texto con tokenizers, procesar datasets grandes usando streaming, y escribir tests para validar tu pipeline de procesamiento de datos.

2. Instalación

pip install datasets transformers

La biblioteca datasets no depende estrictamente de transformers, pero la sección de tokenización en el apartado 6 la necesita.

import datasets
print(datasets.__version__)

3.0.1

3. Cargar datasets desde Hugging Face Hub

3.1. load_dataset básico

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 carga el script/metadatos del dataset imdb desde el Hub, almacena en caché el archivo comprimido localmente (~/.cache/huggingface/datasets), y devuelve un DatasetDict: un diccionario de Dataset organizados por split (train, test, unsupervised).

3.2. Cargar un split específico

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

<class 'datasets.arrow_dataset.Dataset'> 25000

Pasar split="train" devuelve directamente un Dataset en lugar de un DatasetDict, lo que es conveniente cuando solo necesitas un split.

Nota: Si el dataset en el Hub requiere una configuración específica (por ejemplo, glue), debes pasar el nombre de la configuración como argumento adicional: load_dataset("glue", "mrpc"). Si omites este parámetro, obtendrás un error pidiendo que selecciones una configuración.

4. Explorar la estructura del 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...

Acceder por índice devuelve un diccionario Python normal. Dataset también admite slicing que devuelve un diccionario de listas:

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

[0, 0, 0]

Ver el esquema usando features:

print(train_reviews.features)

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

ClassLabel te indica que 0 corresponde a la etiqueta neg y 1 a pos: útil cuando necesites mapear de vuelta a los nombres de las etiquetas.

5. Procesar datos con map, filter, select

5.1. filter para filtrar filas

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

12500

filter acepta una función que devuelve True/False para cada fila, itera sobre todo el dataset y solo conserva las filas que cumplan la condición.

5.2. map para transformar datos

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 aplica una función a cada fila y devuelve un diccionario con campos nuevos (u campos existentes actualizados): los resultados se combinan nuevamente con el dataset original.

5.3. map en lotes para acelerar

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)

Con batched=True, la función recibe un diccionario de lote completo (cada campo es una lista) en lugar de una fila por vez, lo que evita llamar a la función Python 25.000 veces: mucho más rápido con datasets grandes.

5.4. select para obtener un subconjunto

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

100

Úsalo cuando quieras probar tu pipeline rápidamente en un subconjunto antes de ejecutar el dataset completo.

6. Preprocesar texto con tokenizer

El caso de uso más común en la práctica: tokenizar el texto para pasarlo al modelo.

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 recibe un lote y devuelve un diccionario con input_ids y attention_mask: map combinará automáticamente estos nuevos campos con el dataset, conservando los campos existentes. Para eliminar campos que no necesites (por ejemplo, el text sin procesar para ahorrar RAM), usa remove_columns:

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

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

7. Dividir splits, barajar, shard

7.1. Dividir train/test manualmente

Si el dataset solo tiene un split, puedes dividirlo tú mismo usando 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 garantiza que el resultado de la división sea reproducible entre ejecuciones.

7.2. Barajar y shard

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

6250

shard divide el dataset en num_shards partes iguales: útil cuando distribuyes datos a múltiples workers en entrenamiento distribuido.

8. Streaming para datasets grandes

Cuando el dataset es demasiado grande para cargarlo completamente en disco/RAM, usa streaming=True:

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

3

Con streaming, datasets no carga el archivo completo sino que lee registros secuencialmente a través de un iterador. No puedes usar len() ni indexación directa, pero aún puedes usar map/filter de forma lazy (se aplican solo cuando iteras).

Nota: Streaming es apropiado para datasets de cientos de GB como los conjuntos de pretraining. Con datasets pequeños (pocos cientos de MB), la carga normal (streaming=False) es más rápida porque aprovecha el caché y el acceso aleatorio.

9. Guardar y recargar datasets

Después de terminar el procesamiento, guarda los resultados para no tener que ejecutar map cada vez:

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 escribe el dataset en formato Arrow: muy rápido de releer porque no necesita reprocesarse desde la fuente original. Si necesitas interactuar con pandas, convierte directamente:

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

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

10. Escribir tests para el pipeline de procesamiento de datos

El pipeline de procesamiento de datos debe tener tests para garantizar que la lógica de map/filter no se rompa silenciosamente cuando edites el código. A continuación se muestran tests usando 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 crea un pequeño dataset en memoria sin necesidad de cargar nada del Hub: ideal para tests unitarios. Los dos tests anteriores verifican el comportamiento correcto de map (generando nuevos campos) y filter (filtrando según la condición) sin depender de datasets reales de Internet.

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. Errores comunes y cuándo no usar datasets

• Olvidar batched=True al procesar datasets grandes: map sin batch es lento porque llama a la función Python para cada fila. Con cientos de miles de filas, siempre prioriza batched=True si tu lógica lo permite.
• Usar índices aleatorios en datasets con streaming: streamed[5] lanzará un error porque un dataset con streaming es un iterador que no soporta acceso aleatorio. Solo usa .take(), .skip(), o itera secuencialmente.
• No establecer seed al usar shuffle/train_test_split: los resultados no serán reproducibles entre ejecuciones, lo que dificulta la depuración cuando comparas métricas entre experimentos.
• Caché descontrolado que crece: cada vez que uses map con una función diferente, datasets escribe un nuevo archivo de caché. Si tu máquina tiene poco espacio en disco, usa load_dataset(..., keep_in_memory=True) o limpia periódicamente ~/.cache/huggingface/datasets.
• No usar datasets para datos relacionales complejos (múltiples tablas con joins, transacciones necesarias). Esta biblioteca está optimizada para datos tabulares/textuales planos usados en ML, no es un reemplazo para bases de datos.

12. Conclusión

La biblioteca datasets te proporciona una API limpia para ir desde “cargar un dataset bruto del Hub” a “un dataset tokenizado, listo para entrenar” con solo algunos comandos load_dataset, map y filter. Con datasets pequeños, cárgatelos normalmente y aprovecha el caché Arrow; con datasets enormes, cambia a streaming=True para no tener que descargar todo el contenido en disco. Cuando construyas un pipeline de procesamiento de datos para un proyecto real, no olvides escribir tests para tus funciones map/filter usando Dataset.from_dict: esto te ayuda a detectar temprano si la lógica de preprocesamiento se rompe inesperadamente.