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. La forma más simple es remove_columns, que además es lo que querés casi siempre después de tokenizar — el modelo no acepta columnas de texto:

def tokenize_and_split(examples):
    return tokenizer(examples["review"], truncation=True, max_length=128,
                     return_overflowing_tokens=True)   # una reseña larga → varios pedazos

tokenized = dataset.map(tokenize_and_split, batched=True,
                        remove_columns=dataset["train"].column_names)   # quedan solo las columnas nuevas
print(len(dataset["train"]), len(tokenized["train"]))   # p. ej. 138514 → 206772: más filas que antes

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. Lo mismo vale para el formato que va al modelo: with_format("torch") hace que cada acceso devuelva tensores en vez de listas, sin convertir nada en disco. La diferencia entre las dos variantes es solo de estilo — set_format modifica el objeto en el lugar, with_format devuelve uno nuevo:

ds = tokenized["train"]
print(type(ds[0]["input_ids"]))          # <class 'list'>

ds_torch = ds.with_format("torch")          # nuevo objeto; ds sigue devolviendo listas
print(type(ds_torch[0]["input_ids"]))    # <class 'torch.Tensor'>

ds.set_format("torch", columns=["input_ids", "attention_mask", "label"])   # en el lugar, y solo esas columnas
ds.reset_format()                            # vuelve a listas / Python nativo

Lo que ninguna de las dos hace es tocar los datos: el archivo Arrow es el mismo antes y después. Si querés una conversión real (una columna nueva con otro tipo), eso es un map.

3. Datasets que no entran en RAM: Apache Arrow y Memory-Mapping

La regla empírica clásica de Pandas establece que manipular un dataset requiere entre 5× y 10× su tamaño en memoria RAM. Si intentás cargar un archivo JSON o CSV de 20 GB en Pandas, el proceso colapsará exigiendo entre 100 y 200 GB de memoria física. La librería datasets de Hugging Face erradica este cuello de botella utilizando el estándar Apache Arrow y mapeo directo a memoria virtual del sistema operativo (memory mapping).

Arquitectura en Memoria: Filas de Python vs Columnas de Apache Arrow

¿Por qué Pandas y Python estándar sufren una explosión de memoria mientras Arrow permanece constante?

• Sobrecarga por Objeto en Python (Row-Oriented): En Python estándar, cada fila de una tabla es un diccionario o una lista que contiene punteros a instancias individuales de PyObject. Un simple entero de 64 bits (que a nivel de CPU ocupa 8 bytes) en Python consume al menos 28 bytes de encabezado (tipo de objeto, conteo de referencias). Una cadena de caracteres pequeña añade 50 bytes adicionales más su longitud. Un dataset con 10 millones de filas y 10 columnas instancia 100 millones de objetos dispersos en el heap, provocando una fragmentación severa y destruyendo la localidad de caché del procesador (L1/L2 cache misses superiores al 40%).
• Disposición Columnar de Apache Arrow: Arrow almacena los datos en memoria por columnas contiguas, sin envoltorios de Python. Para una columna de texto, Arrow utiliza únicamente tres buffers contiguos de bytes:
  1. Validity Bitmap: Un arreglo de bits contiguo donde cada bit (0 o 1) indica si el valor es nulo o válido (1 byte cubre 8 filas).
  2. Offsets Buffer: Un buffer contiguo de enteros de 32 o 64 bits que marca el inicio y fin en bytes de cada string dentro del buffer de datos.
  3. Data Buffer: Una secuencia contigua y compacta de bytes con todos los caracteres concatenados.

Zero-Copy IPC y Paginación por Demanda (Demand Paging)

Debido a que el formato de memoria de Arrow está estrictamente estandarizado, datasets implementa Zero-Copy IPC (Inter-Process Communication): transferir datos entre procesos de Python, extensiones de C++, Rust o motores analíticos (como Polars o DuckDB) no requiere serializar ni deserializar datos. Ambos procesos simplemente comparten un puntero a la misma región de memoria compartida sin costo de copia.

Cuando cargás un dataset local de 50 GB desde disco mediante load_from_disk() o load_dataset():

1. El archivo .arrow en disco se mapea al espacio de memoria virtual del proceso con mmap(MAP_SHARED).
2. La memoria residente (RSS) del proceso de Python permanece prácticamente en cero bytes al iniciar.
3. Cuando el bucle de entrenamiento o el DataLoader solicita el batch i, la CPU genera Page Faults de hardware y el kernel de Linux transfiere únicamente las páginas físicas de 4 KB correspondientes a ese lote desde el disco al Page Cache de la RAM.
4. Una vez procesadas y copiadas a la VRAM de la GPU, si el sistema sufre presión de memoria, el kernel puede descartar inmediatamente esas páginas de RAM sin escribirlas a disco (ya que son de solo lectura).

Streaming Masivo con IterableDataset

Cuando el corpus supera la capacidad del almacenamiento local (por ejemplo, datasets de preentrenamiento como FineWeb o RedPajama, de decenas de terabytes), ni siquiera el disco es suficiente. En este escenario, streaming=True conmuta la arquitectura a un IterableDataset:

from datasets import load_dataset

# Descarga diferida (lazy evaluation) por HTTP chunked transfer
dataset = load_dataset("HuggingFaceFW/fineweb", split="train", streaming=True)
print(next(iter(dataset)))  # Lee el primer ejemplo; el disco local no almacena nada

Un IterableDataset opera como un generador de Python sobre conexiones HTTP persistentes. Para mezclar ejemplos sin descargar el archivo completo, IterableDataset.shuffle(buffer_size=10_000) mantiene una ventana deslizante de 10.000 ejemplos en memoria RAM, extrayendo pseudo-aleatoriamente elementos del buffer y reponiéndolos con nuevos datos del flujo entrante.

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.

Y ni siquiera hace falta pasar por un archivo: si los datos ya están en memoria — un DataFrame que armaste en Pandas, o listas paralelas que salieron de un scraper — Dataset.from_pandas y Dataset.from_dict los convierten a Arrow directamente:

from datasets import Dataset

ds = Dataset.from_dict({"text": ["primer issue", "segundo issue"], "n_comments": [3, 0]})
ds = Dataset.from_pandas(df)               # agrega una columna "__index_level_0__" si el índice del df no es el default
ds = Dataset.from_pandas(df, preserve_index=False)   # ...salvo que le digas que no
print(ds.features)                          # los tipos se infieren: Value('string'), Value('int64')...

Es el camino inverso de set_format("pandas"): de Pandas a Arrow, en vez de una vista Arrow con cara de Pandas. Un patrón habitual es hacer la limpieza fea en un DataFrame chico y pasar a Dataset para el map tokenizado y el push_to_hub.

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.

Las dos variantes, en código. El detalle que define la validez matemática del mean pooling es la attention mask:

Formalmente, dado un lote de secuencias donde cada token $i$ tiene un estado oculto de salida ${\mathbf{h}}_i\in {\mathbb{R}}^d$ y una máscara binaria $m_i\in \{0,1\}$ (donde 1 representa texto real y 0 representa relleno [PAD]), el embedding condensado de la oración $\mathbf{e}\in {\mathbb{R}}^d$ se define estrictamente como:

¿Por qué es indispensable este enmascaramiento y por qué no se puede simplemente dividir por la longitud total del tensor $N$?

• Los tokens de padding NO son vectores nulos (${\mathbf{h}}_{\text{pad}}\ne \mathbf{0}$): Debido a las capas de Layer Normalization ($\text{LN}(\mathbf{x})=\frac{\mathbf{x}-\mu }{\sigma }\mathit{\gamma }+\mathit{\beta }$) y a los biases de las capas feed-forward en la arquitectura Transformer, las posiciones correspondientes a [PAD] producen vectores con magnitud significativa. Si se promedian sin máscara, los tokens de padding actúan como un atractor artificial que desvía drásticamente el significado semántico del texto.
• Sesgo de Dilución Dependiente de la Longitud: Si se dividiera la suma enmascarada por la longitud máxima del batch $N_{\text{max}}$ (ej. 512) en lugar de la longitud real $N_{\text{real}}=\sum m_i$ (ej. 10 tokens), la representación de una oración corta sufriría una atenuación artificial de más de $50\times$, distorsionando cualquier comparación de distancias o proyecciones.

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = AutoModel.from_pretrained(model_ckpt)

def cls_pooling(model_output):
    return model_output.last_hidden_state[:, 0]          # el vector del primer token ([CLS])

def mean_pooling(model_output, attention_mask):
    tokens = model_output.last_hidden_state                # [batch, seq, hidden]
    mask = attention_mask.unsqueeze(-1).float()             # [batch, seq, 1]: 1 en texto real, 0 en padding
    return (tokens * mask).sum(1) / mask.sum(1).clamp(min=1e-9)

def get_embeddings(textos, pooling="mean"):
    enc = tokenizer(textos, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        out = model(**enc)
    emb = mean_pooling(out, enc["attention_mask"]) if pooling == "mean" else cls_pooling(out)
    return torch.nn.functional.normalize(emb, p=2, dim=1)   # norma 1 → producto interno = coseno

Demostración Matemática: Equivalencia en la Hiperesfera Unitaria ${\mathbb{S}}^{d-1}$

La última línea del pipeline de embeddings (torch.nn.functional.normalize(emb, p=2, dim=1)) es la pieza angular que unifica la geometría del espacio vectorial. FAISS por defecto implementa índices basados en Distancia Euclidiana (${\mathcal{L}}_2$), mientras que los modelos de representación semántica se entrenan maximizando la Similitud Coseno o el Producto Interno.

Veamos la demostración paso a paso de por qué, bajo normalización euclidiana, ambas métricas inducen exactamente el mismo orden topológico:

Sean dos vectores normalizados en la hiperesfera unitaria $\mathbf{u},\mathbf{v}\in {\mathbb{S}}^{d-1}$, de modo que $\parallel \mathbf{u}{\parallel }_2=\parallel \mathbf{v}{\parallel }_2=1$. Desarrollemos el cuadrado de su distancia euclidiana:

Sustituyendo las normas unitarias ($\parallel \mathbf{u}{\parallel }_2^2=1$ y $\parallel \mathbf{v}{\parallel }_2^2=1$):

Despejando la similitud coseno (o producto interno):

Dado que la función $f(x)=1-\frac{1}{2}x$ es estrictamente decreciente y monótona respecto al cuadrado de la distancia, se concluye la equivalencia de ranking:

Consecuencia práctica en FAISS y Bases Vectoriales: El índice por defecto de FAISS (IndexFlatL2) devuelve con exactitud matemática el mismo orden topológico de vecinos más cercanos que un índice de producto interno (IndexFlatIP). Sin normalizar, en cambio, un pasaje largo con norma $\parallel \mathbf{v}{\parallel }_2=10$ superaría espuriamente en producto interno a un pasaje corto de norma $\parallel \mathbf{v}{\parallel }_2=1$, aun cuando el pasaje corto compartiera un ángulo coseno mucho más cercano a la consulta.

Para configurar explícitamente el índice en FAISS con producto interno sobre vectores unitarios:

import faiss

dataset = dataset.map(lambda x: {"embeddings": get_embeddings(x["text"]).numpy()}, batched=True, batch_size=64)
dataset.add_faiss_index(column="embeddings", metric_type=faiss.METRIC_INNER_PRODUCT)

pregunta_emb = get_embeddings(["¿Cómo cargo un dataset offline?"]).numpy()[0]
scores, ejemplos = dataset.get_nearest_examples("embeddings", pregunta_emb, k=5)
# con vectores normalizados, scores son cosenos: 1.0 = idéntico, 0 = nada que ver

Ese es el mismo e5-large + ChromaDB de tu proyecto, en su forma mínima: e5 espera mean pooling y vectores normalizados (y los prefijos query: / passage: en el texto, según su model card), y ChromaDB con hnsw:space = "cosine" hace exactamente lo que hace acá METRIC_INNER_PRODUCT sobre vectores de norma 1.

6. Errores comunes

La mayoría viene de confundir "cambiar los datos" con "cambiar cómo los veo", o de esperar que map adivine.

7. 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; y from_pandas/from_dict ni siquiera necesitan un archivo.
2. map y filter cambian los datos; set_format/with_format cambian solo cómo los ves. Si map cambia el número de filas, remove_columns no es opcional.
3. 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.
4. Memory mapping (Arrow) y streaming resuelven el mismo problema en dos escalas distintas: no cargar todo a RAM, y no descargar todo al disco.
5. Crear un dataset propio es API pública → JSON → load_dataset() → push_to_hub(), con una dataset card documentando qué hay adentro.
6. 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. Mean pooling con attention mask y normalización L2 son las dos líneas que separan "funciona" de "funciona bien".

8. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

9. Para practicar

from datasets import Dataset
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
ds = Dataset.from_dict({
    "text": ["cargar un dataset offline", "el modelo no entra en memoria", "cómo subo un dataset al hub", None],
    "label": [0, 1, 0, 1],
})
print(ds.features)

ds = ds.filter(lambda x: x["text"] is not None)   # sacá esta línea y mirá qué error da el map
tokenized = ds.map(lambda x: tok(x["text"], truncation=True), batched=True, remove_columns=["text"])
print(tokenized.column_names, type(tokenized[0]["input_ids"]))

t = tokenized.with_format("torch")
print(type(t[0]["input_ids"]), type(tokenized[0]["input_ids"]))   # tensor vs. list: el original no cambió
print(t.cache_files == tokenized.cache_files)                # mismo archivo Arrow detrás

import torch
from transformers import AutoTokenizer, AutoModel

ckpt = "sentence-transformers/all-MiniLM-L6-v2"
tok, model = AutoTokenizer.from_pretrained(ckpt), AutoModel.from_pretrained(ckpt)

def embed(textos):
    enc = tok(textos, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        h = model(**enc).last_hidden_state
    mask = enc["attention_mask"].unsqueeze(-1).float()
    return (h * mask).sum(1) / mask.sum(1)

frases = ["How do I load a dataset without internet?", "Loading data offline",
          "The model runs out of GPU memory", "My cat sleeps all day"]
e = embed(frases)
print(e.norm(dim=1))                                   # normas distintas de 1 y distintas entre sí
en = torch.nn.functional.normalize(e, p=2, dim=1)
print((en @ en.T).round(decimals=2))                   # producto interno de normalizados
print(torch.nn.functional.cosine_similarity(e[0], e[1], dim=0))   # coseno de los crudos: mismo número

import faiss
from datasets import Dataset

ds = Dataset.from_dict({"text": frases})
ds = ds.map(lambda x: {"emb": torch.nn.functional.normalize(embed(x["text"]), dim=1).numpy()}, batched=True)
ds.add_faiss_index(column="emb", metric_type=faiss.METRIC_INNER_PRODUCT)

q = torch.nn.functional.normalize(embed(["reading files with no network connection"]), dim=1).numpy()[0]
scores, ej = ds.get_nearest_examples("emb", q, k=2)
for s, t in zip(scores, ej["text"]):
    print(f"{s:.2f}  {t}")

10. Fuentes y para seguir

• Hugging Face LLM Course — Chapter 5: The 🤗 Datasets library — el capítulo original, con el ejemplo completo de los issues de GitHub y el buscador con FAISS.
• Datasets — documentación — la referencia de map, filter, formatos, streaming y add_faiss_index; la guía "Process" es la que más vas a abrir.
• Apache Arrow — el formato columnar detrás de todo el capítulo; leé la sección de memory mapping para entender §3 de verdad.
• FAISS (Facebook AI Similarity Search) — la librería de índices vectoriales; su wiki explica cuándo un índice plano deja de alcanzar y qué usar después (IVF, HNSW).
• Sentence-BERT (Reimers & Gurevych, 2019) — el paper que estableció mean pooling + coseno como la receta estándar de embeddings de oración.
• Text Embeddings by Weakly-Supervised Contrastive Pre-training (Wang et al., 2022) — el paper de E5, la familia del multilingual-e5-large de tu RAG; explica los prefijos query: / passage:.