Skip to content

English · Español

Lab 00 — Construir el batcher y la mask

Objetivo: ensamblar el pipeline de datos que convierte los shards JSONL de conjugaciones verbales de la Fase 12 en cuádruples (input, target, loss_mask, attn_pad_mask) en batches, listos para alimentar a MiniGPT.

Tiempo estimado: 90–120 minutos.

Requisito previo: encoder BPE de la Fase 11 + corpus JSONL de la Fase 12 commiteados.

Lo que produces

Un nuevo archivo de módulo + tests:

  • src/minitrain/data.py — cargador de corpus, shuffle determinista, batcher, constructor de mask, ayudante de split held-out.
  • tests/minitrain/test_data.py — tests de esquema, formas, mask y split.

Un pequeño experimento:

  • experiments/18-batch-sanity/ — manifest.json + un batch de ejemplo impreso con formas y una frase decodificada por fila.

Antecedentes que debes haber leído

  • theory/01-batching-loss-mask.md — las dos masks, el desplazamiento causal, reducción token-mean, política de split held-out.
  • corpus_spec.md de la Fase 12 — esquema JSONL para entradas de conjugación verbal.

TODOs

Bloque A — cargador de corpus

Implementa:

@dataclass(frozen=True)
class VerbExample:
    id: str                # e.g. "work.present.1sg"
    en: str                # "I work"
    es: str                # "yo trabajo"
    verb: str              # "work"
    tense: str             # "present_simple"
    person: str            # "1sg"
    is_regular: bool

def load_corpus(path: Path) -> list[VerbExample]: ...
  • Parsea data/processed/{train,val,test}.jsonl (salida de la Fase 12).
  • Valida: toda fila tiene todos los campos requeridos; verb ∈ {20 verbos}, tense ∈ {5 tiempos}, person ∈ {1sg, 2sg, 3sg}.
  • Rechaza valores desconocidos con un error claro citando el id de la fila ofensora.

Bloque B — constructor de ejemplo tokenizado

El modelo entrena con secuencias tokenizadas de la forma:

<bos> es:<spanish_form> / en:<english_form> <eos>

(Ambas direcciones están presentes en el corpus — el modelo aprende el emparejamiento.)

  • Implementa tokenize_example(ex: VerbExample, tokenizer) -> list[int] devolviendo los ids de tokens. Usa el BPE de la Fase 11.
  • Añade los special tokens <bos>, <eos>, <pad>; sus ids deben estar reservados en el vocabulario de la Fase 11. Si no, extiende el vocabulario de forma determinista (el lab de la Fase 11 fijó esto).
  • Test de propiedad: decode(tokenize_example(ex)) == "<bos> es:... / en:... <eos>".

Bloque C — split held-out

Implementa la política hold-out-4-verbs de theory/01:

HELDOUT_VERBS = ("look", "like", "see", "eat")  # 2 regular + 2 irregular

def held_out_split(examples: list[VerbExample]) -> tuple[list, list]:
    """Returns (train, val). Val is all examples whose verb is in HELDOUT_VERBS."""
    ...
  • Train debería contener 16 verbos × 5 tiempos × 3 personas = 240 formas.
  • Val debería contener 4 verbos × 5 tiempos × 3 personas = 60 formas.
  • Sin fugas: train ∩ val == ∅ (por id de ejemplo).
  • Test: asegúrate de los tamaños; comprueba que los verbos held-out son exactamente los de HELDOUT_VERBS; comprueba que los 5 tiempos y las 3 personas aparecen en ambos train y val.

Bloque D — batcher y masks

Implementa:

def make_batch(
    examples: list[list[int]],
    pad_id: int,
) -> tuple[ndarray, ndarray, ndarray, ndarray]:
    """Returns (input, target, loss_mask, attn_pad_mask)."""

Las formas tras el desplazamiento causal son todas (B, L-1). Consulta el esbozo de implementación canónico en theory/01 — puedes transcribirlo, pero debes añadir los tests del Bloque E.

  • input.shape == target.shape == loss_mask.shape == attn_pad_mask.shape == (B, L-1).
  • loss_mask es 1 donde target != pad_id, si no 0.
  • attn_pad_mask es 1 donde input != pad_id, si no 0.
  • input y target son int32; loss_mask y attn_pad_mask son float32.

Bloque E — shuffle determinista

Implementa un BatchIterator:

class BatchIterator:
    def __init__(self, examples: list[list[int]], batch_size: int, seed: int): ...
    def __iter__(self) -> Iterator[tuple[ndarray, ndarray, ndarray, ndarray]]: ...
    def state_dict(self) -> dict: ...
    def load_state_dict(self, state: dict) -> None: ...
  • Cada época hace shuffle de los ejemplos con un RNG sembrado, distinto del RNG de control de entrenamiento.
  • state_dict() devuelve la época actual, la posición y el estado del RNG; el round-trip vía load_state_dict reanuda la misma secuencia.
  • El último batch de cada época puede ser menor que batch_size; no lo descartes.

Bloque F — tests

En tests/minitrain/test_data.py:

  1. test_corpus_schema — load_corpus(train.jsonl) valida sin error; conteo esperado = 240.
  2. test_heldout_split_sizes — split 240/60, disjunto por id.
  3. test_heldout_coverage — train y val contienen los 5 tiempos y las 3 personas; train tiene 16 verbos, val tiene 4.
  4. test_batch_shapes — dados ejemplos de longitudes [6, 11, 9, 14], las formas del batch son (4, 13) (tras padding a 14 y desplazamiento causal a 13).
  5. test_batch_masks — comprueba que loss_mask tiene exactamente el conteo correcto de ceros (las posiciones de pad en los targets).
  6. test_iterator_resume — itera 5 batches, captura state_dict, itera 5 más; recarga desde el snapshot, itera 5 más, verifica que los batches son los mismos que los segundos-5 del original.
  7. test_no_overlap_train_val — todo id de ejemplo en val no está en train.

Bloque G — experimento de sanity

experiments/18-batch-sanity/:

  • manifest.json con seed, versions, config (batch_size, hash del tokenizer).
  • print_sample.py — carga el corpus, construye un batch de tamaño 4, decodifica cada fila de vuelta a texto, imprime formas y texto decodificado.
  • Salida: sample_batch.txt con las 4 frases decodificadas y los 4 vectores de mask como cadenas [1 1 1 1 1 0 0 0 ...].
  • Inspección visual: los ceros de la mask se alinean con los tokens de pad en las frases decodificadas.

Restricciones

  • NumPy puro + librería estándar.
  • Sin Pandas (excesivo para 600 ejemplos).
  • El RNG del shuffle debe ser un np.random.Generator nuevo, no el RNG global, para que no se enrede con el RNG de control de entrenamiento.

Condiciones de parada

Hecho cuando:

  1. pytest tests/minitrain/test_data.py -v pasa los siete tests.
  2. experiments/18-batch-sanity/sample_batch.txt está commiteado con frases decodificadas y vectores de mask.
  3. Los números del split held-out coinciden con theory/01 (240 train / 60 val).
  4. Puedes reformular la distinción entre loss-mask y attention-mask sin consultar el archivo.

Escollos

  • Poner <pad> antes de <eos> en una fila. No hagas pre-pad; haz right-pad después de <eos>. Si tu tokenizer maneja mal esto, extiéndelo.
  • Shuffle aleatorio usando el RNG global. Esto rompe la reproducibilidad al recargar-y-reanudar. Usa un Generator dedicado.
  • Asignar un buffer nuevo en cada batch. Para 600 ejemplos está bien; para corpora grandes lo reutilizarías. Anótalo en un comentario y sigue.
  • Olvidar que el target es input[:, 1:], no input[:, :-1]. La mask se deriva del target, no del input.

Cuándo consultar solutions/

Después de que los siete tests pasen. Solución en solutions/00-batch-and-mask-ref.md (escrita al abrir la fase).

Siguiente lab: lab/01-train-mini.md.