La librería Datasets: de un CSV local a un buscador semántico

1. Cuando el dataset no está en el Hub

load_dataset() no está atado al Hub — el primer argumento puede ser un tipo de archivo en vez de un nombre de dataset:

data_files acepta un path local, una URL remota, o un diccionario que mapea nombres de split a archivos — y descomprime .gz/.zip/.tar automáticamente:

from datasets import load_dataset

data_files = {"train": "train.json.gz", "test": "test.json.gz"}
dataset = load_dataset("json", data_files=data_files, field="data")

Cargar desde una URL es literalmente lo mismo — data_files no distingue entre "local" y "remoto", así que todo lo que aprendas acá aplica sin cambios a un archivo en tu disco o en un bucket.

2. Cortar y trocear

Para tener una primera impresión de un dataset grande, combiná shuffle() y select() para tomar una muestra chica; usá filter() (con una lambda) para sacar filas problemáticas antes de mapear una función que no tolera esos casos:

muestra = dataset["train"].shuffle(seed=42).select(range(1000))

# .map() sobre un valor None revienta — filtrar primero
dataset = dataset.filter(lambda x: x["condition"] is not None)
dataset = dataset.map(lambda x: {"condition": x["condition"].lower()})

El argumento batched=True de map() no es un detalle menor — es la diferencia entre minutos y segundos. Los números reales del curso, tokenizando el mismo corpus:

Un tokenizer rápido con batched=True es ~30× más rápido que uno lento sin batchear — la parte en Rust de tokenizers paraleliza internamente, algo que num_proc (multiprocessing de Python) no puede igualar para el tokenizer rápido, pero sí ayuda mucho al lento.

Cuando una función de map() cambia el número de filas (por ejemplo, tokenizar con return_overflowing_tokens=True parte una reseña larga en varios ejemplos), hay que decirle explícitamente qué columnas viejas descartar o cómo remapearlas — si no, el error es un ArrowInvalid por longitudes que no coinciden.

Para EDA con Pandas, set_format("pandas") cambia solo el formato de salida, no el dato subyacente (que sigue siendo Arrow) — y train_test_split() te da un split de validación con una línea, con semilla para reproducibilidad.

3. Datasets que no entran en RAM

La regla de Pandas dice que necesitás 5-10× la memoria del tamaño del dataset. datasets no sigue esa regla porque no carga el dataset a RAM — lo trata como un archivo memory-mapped en formato Apache Arrow: podés operar sobre 20 GB de datos con una fracción de esa RAM, porque el sistema operativo trae a memoria solo las páginas que realmente tocás.

Cuando ni siquiera el disco alcanza (el ejemplo del original: 825 GB), streaming=True cambia el tipo de retorno a IterableDataset — descarga y procesa un ejemplo a la vez, sin bajar el dataset completo:

dataset = load_dataset("json", data_files=url, split="train", streaming=True)
next(iter(dataset))  # primer ejemplo, sin haber descargado el resto

IterableDataset.shuffle(buffer_size=...) mezcla dentro de un buffer, no el dataset entero (no puede — no lo tiene completo en ningún lado), e interleave_datasets() combina varias fuentes streameadas en una sola, alternando entre ellas.

4. Crear tu propio dataset

El ejemplo del original — scrapear los issues de GitHub del propio repo de datasets vía la REST API, filtrar pull requests, sumar los comentarios de cada issue — sigue siendo el mejor ejemplo didáctico de un patrón real: API pública → JSON crudo → load_dataset("json", ...). No hay nada especial en el formato final; cualquier fuente que puedas convertir a JSON Lines entra por la misma puerta que usaste en la sección 1.

Una vez armado, subirlo es la misma llamada que ya conocés del Capítulo 4 — push_to_hub() también existe en Dataset, no solo en modelos:

dataset.push_to_hub("mi-usuario/github-issues")

Y como con los modelos, un dataset sin documentar es un dataset a medio publicar — la dataset card (mismo README.md, mismo principio) es lo que le permite a otra persona (o a vos, en seis meses) saber qué hay adentro sin tener que inspeccionar filas a mano.

5. Búsqueda semántica con FAISS

Esta es la sección que más vale la pena mirar de cerca, porque es — a mano, en unas pocas líneas — exactamente lo que hace un sistema de RAG: convertir texto a vectores, e indexarlos para buscar por significado en vez de por palabra exacta.

1. Embeber cada documento: un modelo encoder-only (Capítulo 1) convierte cada texto en un vector, agrupando el estado oculto del token [CLS] como resumen de toda la secuencia.
2. Indexar los vectores: Dataset.add_faiss_index(column="embeddings") arma una estructura de búsqueda por similitud sobre esa columna.
3. Embeber la pregunta con el mismo modelo, y pedirle al índice los vecinos más cercanos:

pregunta_emb = get_embeddings(["¿Cómo cargo un dataset offline?"])
scores, ejemplos = dataset.get_nearest_examples("embeddings", pregunta_emb, k=5)

Ni una palabra de la pregunta necesita aparecer literalmente en la respuesta — lo que se compara son vectores de significado, no substrings. Esa es la diferencia de fondo entre búsqueda semántica y Ctrl+F.

Un detalle de pooling que vale la pena señalar porque cambia según el modelo: el original usa CLS pooling (tomar el vector del token especial [CLS]), pero muchos modelos de embeddings modernos —incluido multilingual-e5-large— promedian los vectores de todos los tokens en vez de quedarse con uno solo. Cuál usar no es una preferencia estética: hay que usar el que el modelo específico espera, documentado en su model card.

6. Resumen

1. load_dataset() no depende del Hub — un tipo de archivo (csv/json/text) más data_files cubre cualquier fuente local o remota.
2. map(batched=True) es la optimización más importante de este capítulo — hasta 30× más rápido combinado con un tokenizer rápido.
3. Memory mapping (Arrow) y streaming resuelven el mismo problema en dos escalas distintas: no cargar todo a RAM, y no descargar todo al disco.
4. Crear un dataset propio es API pública → JSON → load_dataset() → push_to_hub(), con una dataset card documentando qué hay adentro.
5. Embeddings + un índice de similitud es búsqueda semántica — la misma idea detrás de cualquier sistema de RAG, incluido el que ya corre en tu propio proyecto.