Reutilizar modelos: BERT, codificador-decodificador y ajuste fino

import numpy as np
import matplotlib.pyplot as plt
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

NOMBRE = "distilbert/distilbert-base-uncased"
tokenizador = AutoTokenizer.from_pretrained(NOMBRE)
bert_mlm = AutoModelForMaskedLM.from_pretrained(NOMBRE).to(device).eval()


def completar(frase, k=5):
    """Las k palabras más probables para el [MASK] de la frase."""
    x = tokenizador(frase, return_tensors="pt").to(device)
    with torch.no_grad():
        logits = bert_mlm(**x).logits
    # La posición del [MASK] dentro de la frase tokenizada
    pos = (x["input_ids"][0] == tokenizador.mask_token_id).nonzero()[0, 0]
    probs = logits[0, pos].softmax(dim=-1)
    mejores = probs.topk(k)
    palabras = tokenizador.convert_ids_to_tokens(mejores.indices.tolist())
    return palabras, mejores.values.tolist()


print(tokenizador.convert_ids_to_tokens(
    tokenizador("I loved this book.")["input_ids"]))
for frase in ["I loved this book, it was [MASK].",
              "The cat sat on the [MASK].",
              "The capital of France is [MASK]."]:
    palabras, probs = completar(frase)
    print(frase)
    print(" ", ", ".join(f"{p} {q:.2f}" for p, q in zip(palabras, probs)))

palabras, probs = completar("I loved this book, it was [MASK].")
colores = ["C1"] + ["C0"] * (len(palabras) - 1)
plt.barh(palabras[::-1], probs[::-1], color=colores[::-1])
plt.xlabel("Probabilidad")
plt.title("«I loved this book, it was ___.»")
plt.show()

import textwrap
from huggingface_hub import hf_hub_download
from datasets import load_dataset

ruta = hf_hub_download("fancyzhx/amazon_polarity",
                       "amazon_polarity/test-00000-of-00001.parquet",
                       repo_type="dataset")
resenas = load_dataset("parquet", data_files=ruta, split="train")
muestra = resenas.shuffle(seed=0).select(range(1900))

# Título y texto de cada reseña juntos; la etiqueta 1 es positiva
datos = [(f"{r['title']}. {r['content']}", r["label"]) for r in muestra]
entrenamiento, validacion = datos[:1500], datos[1500:]

print(len(resenas), "reseñas en el archivo")
for nombre, parte in [("Entrenamiento", entrenamiento),
                      ("Validación", validacion)]:
    positivas = sum(etiqueta for _, etiqueta in parte)
    print(f"{nombre}: {len(parte)} reseñas, {positivas} positivas")
texto, etiqueta = entrenamiento[0]
print(textwrap.fill(texto[:230] + "…", 76), "→", etiqueta)

from torch.utils.data import DataLoader


def juntar(lote):
    """Tokeniza una lista de (texto, etiqueta) y la rellena al más largo."""
    textos = [texto for texto, _ in lote]
    x = tokenizador(textos, truncation=True, max_length=128,
                    padding=True, return_tensors="pt")
    x["labels"] = torch.tensor([etiqueta for _, etiqueta in lote])
    return x


def cargador_entrenamiento():
    # El mismo orden de lotes cada vez, para comparar en igualdad
    return DataLoader(entrenamiento, batch_size=16, shuffle=True,
                      collate_fn=juntar,
                      generator=torch.Generator().manual_seed(0))


cargador_val = DataLoader(validacion, batch_size=32, collate_fn=juntar)


def ajustar(modelo, tasa):
    """Una época del bucle; devuelve la pérdida de cada lote."""
    modelo.to(device)
    # Solo se optimizan los parámetros que no están congelados
    entrenables = [p for p in modelo.parameters() if p.requires_grad]
    optimizador = torch.optim.AdamW(entrenables, lr=tasa)
    modelo.train()
    perdidas = []
    for lote in cargador_entrenamiento():
        lote = lote.to(device)
        optimizador.zero_grad()
        salida = modelo(**lote)          # con labels, trae la pérdida
        salida.loss.backward()
        optimizador.step()
        perdidas.append(salida.loss.item())
    return perdidas


def medir(modelo):
    """Pérdida media y exactitud en el conjunto de validación."""
    modelo.eval()
    suma, aciertos = 0.0, 0
    with torch.no_grad():
        for lote in cargador_val:
            lote = lote.to(device)
            salida = modelo(**lote)
            etiquetas = lote["labels"]
            suma += salida.loss.item() * len(etiquetas)
            aciertos += (salida.logits.argmax(dim=1) == etiquetas).sum()
    return suma / len(validacion), aciertos.item() / len(validacion)

from transformers import AutoModelForSequenceClassification


def nuevo_clasificador():
    """DistilBERT preentrenado con una cabeza nueva de 2 clases."""
    torch.manual_seed(0)            # la misma cabeza inicial cada vez
    return AutoModelForSequenceClassification.from_pretrained(
        NOMBRE, num_labels=2)


def contar(modelo):
    entrenables = sum(p.numel() for p in modelo.parameters()
                      if p.requires_grad)
    total = sum(p.numel() for p in modelo.parameters())
    return entrenables, total


resultados = {}


def probar(nombre, modelo, tasa):
    perdidas = ajustar(modelo, tasa)
    perdida_val, exactitud = medir(modelo)
    entrenables, total = contar(modelo)
    resultados[nombre] = {"entrenables": entrenables,
                          "perdidas": perdidas, "exactitud": exactitud}
    print(f"{nombre}: {entrenables:,} de {total:,} entrenables"
          .replace(",", " "))
    print(f"  validación: pérdida {perdida_val:.3f}, "
          f"exactitud {exactitud:.3f}")

cabeza = nuevo_clasificador()
for p in cabeza.distilbert.parameters():    # el codificador, congelado
    p.requires_grad = False
probar("Solo la cabeza", cabeza, tasa=1e-3)

completo = nuevo_clasificador()
probar("Completo", completo, tasa=5e-5)

from matplotlib.patches import Rectangle

fig, ax = plt.subplots(figsize=(4.8, 2.6))
# Dibujo esquemático: r está exagerado para que se vea
r_dibujo = 0.16
piezas = [(0.0, 0.0, 1.0, 1.0, "C3", "W\n768 × 768\ncongelada"),
          (1.45, 0.0, r_dibujo, 1.0, "C0", "B"),
          (1.85, 1.0 - r_dibujo, 1.0, r_dibujo, "C2", "A")]
for x, y, ancho, alto, color, texto in piezas:
    ax.add_patch(Rectangle((x, y), ancho, alto, color=color))
ax.text(0.5, 0.5, piezas[0][5], ha="center", va="center", fontsize=13)
ax.text(1.45 + r_dibujo / 2, -0.1, "B\n768 × r", ha="center", va="top",
        fontsize=13)
ax.text(2.35, 0.75, "A\nr × 768", ha="center", va="top", fontsize=13)
ax.text(1.22, 0.5, "+", ha="center", va="center", fontsize=22)
ax.text(1.73, 0.5, "·", ha="center", va="center", fontsize=34)
ax.set_xlim(-0.05, 2.9)
ax.set_ylim(-0.45, 1.05)
ax.set_aspect("equal")
ax.axis("off")
ax.set_title("LoRA: W + B·A (solo B y A se entrenan)")
plt.show()

from peft import LoraConfig, get_peft_model

config = LoraConfig(task_type="SEQ_CLS", r=8, lora_alpha=16,
                    lora_dropout=0.1, target_modules=["q_lin", "v_lin"])
lora = get_peft_model(nuevo_clasificador(), config)

# Cuántos números agrega LoRA: 6 capas, 2 matrices, B y A
solo_lora = sum(p.numel() for n, p in lora.named_parameters()
                if "lora_" in n)
print("Parámetros de LoRA:", solo_lora, "=", 6 * 2 * (768 * 8 + 8 * 768))
probar("LoRA (r = 8)", lora, tasa=5e-4)

nombres = list(resultados)
plt.figure()
for nombre, color in zip(nombres, ["C2", "C0", "C1"]):
    perdidas = np.array(resultados[nombre]["perdidas"])
    # Promedio móvil de 10 lotes, para ver la tendencia
    suave = np.convolve(perdidas, np.ones(10) / 10, mode="valid")
    plt.plot(np.arange(10, len(perdidas) + 1), suave, color=color,
             label=nombre)
plt.xlabel("Lote")
plt.ylabel("Pérdida de entrenamiento")
plt.title("Una época sobre 1500 reseñas")
# Espacio arriba para la leyenda
plt.ylim(0.2, 0.95)
plt.legend(loc="upper right")
plt.show()

exactitudes = [resultados[n]["exactitud"] for n in nombres]
entrenables = [resultados[n]["entrenables"] for n in nombres]
plt.figure()
plt.barh(nombres, exactitudes, color=["C0", "C0", "C1"])
for i, (e, n) in enumerate(zip(exactitudes, entrenables)):
    texto = f"{n / 1e6:.2f} M" if n < 1e7 else f"{n / 1e6:.0f} M"
    plt.text(0.03, i, f"{e:.1%}  ·  {texto} entrenables", va="center",
             color="white", fontsize=13)
plt.xlim(0, 1)
plt.xlabel("Exactitud en validación")
plt.title("Exactitud según la forma")
plt.show()

def probabilidad_positiva(modelo, textos):
    modelo.eval()
    x = tokenizador(textos, truncation=True, max_length=128,
                    padding=True, return_tensors="pt").to(device)
    with torch.no_grad():
        probs = modelo(**x).logits.softmax(dim=-1)[:, 1]
    return probs.tolist()


propias = ["Not bad at all, it works better than I expected.",
           "I returned it the next day.",
           "Great price, but it broke after a week."]
for texto, p in zip(propias, probabilidad_positiva(lora, propias)):
    print(f"{p:.2f}  {texto}")

from transformers import AutoModelForSeq2SeqLM

tok_en_es = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
# "eager": la atención clásica, que puede devolver sus pesos
en_es = AutoModelForSeq2SeqLM.from_pretrained(
    "Helsinki-NLP/opus-mt-en-es", attn_implementation="eager"
).to(device).eval()


def traducir(frases, tok, modelo):
    x = tok(frases, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        ids = modelo.generate(**x, do_sample=False, num_beams=1,
                              max_new_tokens=60)
    return tok.batch_decode(ids, skip_special_tokens=True)


frases = ["The black cat sleeps on the old chair.",
          "Not bad at all, it works better than I expected.",
          "I returned it the next day."]
for original, traduccion in zip(frases, traducir(frases, tok_en_es,
                                                 en_es)):
    print(original)
    print("  →", traduccion)

from matplotlib.colors import LinearSegmentedColormap

x = tok_en_es(frases[0], return_tensors="pt").to(device)
with torch.no_grad():
    ids = en_es.generate(**x, do_sample=False, num_beams=1)
    salida = en_es(**x, decoder_input_ids=ids[:, :-1],
                   output_attentions=True)
# Cuarta capa de 6, promedio de sus 8 cabezas: (salida, entrada)
pesos = salida.cross_attentions[3][0].mean(dim=0).cpu().numpy()
# En Marian, «▁» marca el comienzo de una palabra, como «Ġ» en Qwen
entrada = [t.replace("▁", "") for t in
           tok_en_es.convert_ids_to_tokens(x["input_ids"][0])]
generada = [t.replace("▁", "") for t in
            tok_en_es.convert_ids_to_tokens(ids[0, 1:])]
print("Entrada:", " ".join(entrada))
print("Salida: ", " ".join(generada))

tinta = LinearSegmentedColormap.from_list("tinta", ["white", "C0"])
plt.figure(figsize=(4.8, 4.8))
plt.imshow(pesos[:-1], cmap=tinta)     # sin la fila de </s>
plt.xticks(range(len(entrada)), entrada, rotation=90)
plt.yticks(range(len(generada) - 1), generada[:-1])
plt.xlabel("Entrada (inglés)")
plt.ylabel("Token que escribe (español)")
plt.title("Atención cruzada")
plt.grid(False)
plt.show()

from PIL import Image
from sklearn.datasets import load_sample_image
from transformers import AutoProcessor, AutoModelForImageTextToText

FLORENCE = "florence-community/Florence-2-base-ft"
procesador = AutoProcessor.from_pretrained(FLORENCE)
descriptor = AutoModelForImageTextToText.from_pretrained(FLORENCE)
descriptor = descriptor.to(device).eval()

fotos = {n: load_sample_image(n) for n in ["china.jpg", "flower.jpg"]}
descripciones = []
for nombre, foto in fotos.items():
    x = procesador(text="<CAPTION>", images=Image.fromarray(foto),
                   return_tensors="pt").to(device)
    with torch.no_grad():
        ids = descriptor.generate(**x, do_sample=False, num_beams=3,
                                  max_new_tokens=30)
    texto = procesador.batch_decode(ids, skip_special_tokens=True)[0]
    descripciones.append(texto)
    print(f"{nombre}: {texto}")

# Y en español, con el traductor de la sección anterior
for texto in traducir(descripciones, tok_en_es, en_es):
    print("→", texto)

fig, ejes = plt.subplots(2, 1, figsize=(4.8, 4.8))
for eje, (nombre, foto) in zip(ejes, fotos.items()):
    eje.imshow(foto)
    eje.set_title(nombre, fontsize=14)
    eje.axis("off")
plt.show()

import torch
from torch import nn


def evaluar(modelo, cargador, funcion_perdida, dispositivo="cpu"):
    """Pérdida media y exactitud de `modelo` sobre `cargador`."""
    modelo.eval()
    suma, aciertos, total = 0.0, 0, 0
    with torch.no_grad():
        for x, y in cargador:
            x, y = x.to(dispositivo), y.to(dispositivo)
            salida = modelo(x)
            suma += funcion_perdida(salida, y).item() * len(y)
            aciertos += (salida.argmax(dim=1) == y).sum().item()
            total += len(y)
    return suma / total, aciertos / total


def entrenar(modelo, cargador_ent, cargador_val, epocas, optimizador,
             funcion_perdida=None, planificador=None, dispositivo="cpu",
             mostrar=True):
    """Entrena `modelo` y devuelve el historial de cada época."""
    if funcion_perdida is None:
        funcion_perdida = nn.CrossEntropyLoss()
    modelo.to(dispositivo)
    historial = {"perdida_ent": [], "exactitud_ent": [],
                 "perdida_val": [], "exactitud_val": []}
    for epoca in range(1, epocas + 1):
        modelo.train()
        suma, aciertos, total = 0.0, 0, 0
        for x, y in cargador_ent:
            x, y = x.to(dispositivo), y.to(dispositivo)
            optimizador.zero_grad()
            salida = modelo(x)
            perdida = funcion_perdida(salida, y)
            perdida.backward()
            optimizador.step()
            suma += perdida.item() * len(y)
            aciertos += (salida.argmax(dim=1) == y).sum().item()
            total += len(y)
        if planificador is not None:
            planificador.step()
        perdida_val, exactitud_val = evaluar(
            modelo, cargador_val, funcion_perdida, dispositivo)
        historial["perdida_ent"].append(suma / total)
        historial["exactitud_ent"].append(aciertos / total)
        historial["perdida_val"].append(perdida_val)
        historial["exactitud_val"].append(exactitud_val)
        if mostrar:
            print(f"Época {epoca}: pérdida {suma / total:.3f}, "
                  f"validación {perdida_val:.3f}, "
                  f"exactitud de validación {exactitud_val:.3f}")
    return historial

from torch.utils.data import TensorDataset, DataLoader
from torchvision import datasets

moda = datasets.FashionMNIST(root="datos", train=True, download=True)
NOMBRES = ["Polera", "Pantalón", "Chompa", "Vestido", "Abrigo",
           "Sandalia", "Camisa", "Zapatilla", "Bolso", "Botín"]
orden = torch.randperm(len(moda.data),
                       generator=torch.Generator().manual_seed(0))
X_moda = moda.data[orden[:24000]].float().div(255).unsqueeze(1)
y_moda = moda.targets[orden[:24000]]
X_ent, y_ent = X_moda[:20000], y_moda[:20000]
X_val, y_val = X_moda[20000:], y_moda[20000:]


def lotes(X, y, mezclar=False):
    """Un DataLoader de lotes de 128 (mezclados) o de 1000 (en orden)."""
    if mezclar:
        return DataLoader(TensorDataset(X, y), batch_size=128,
                          shuffle=True,
                          generator=torch.Generator().manual_seed(0))
    return DataLoader(TensorDataset(X, y), batch_size=1000)


def red_nueva():
    torch.manual_seed(0)      # los mismos pesos iniciales cada vez
    return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
                         nn.Linear(128, 10))


red = red_nueva()
optimizador = torch.optim.Adam(red.parameters(), lr=1e-3)
historial = entrenar(red, lotes(X_ent, y_ent, mezclar=True),
                     lotes(X_val, y_val), 3, optimizador)

azar = torch.Generator().manual_seed(1)
ruido = torch.randn(X_val.shape, generator=azar)
dominios = {
    "Normal": X_val,
    # rot90 gira un cuarto de vuelta en los ejes 2 y 3 (alto y ancho)
    "Rotada": torch.rot90(X_val, 1, dims=(2, 3)),
    "Oscura": 0.4 * X_val,
    # clamp recorta: lo que pasa de 1 queda en 1 y lo negativo, en 0
    "Con ruido": (X_val + 0.3 * ruido).clamp(0, 1),
}
funcion_perdida = nn.CrossEntropyLoss()
acierto_dominio = {}
for nombre, X_d in dominios.items():
    _, acierto_dominio[nombre] = evaluar(red, lotes(X_d, y_val),
                                         funcion_perdida)
    print(f"{nombre:10} exactitud {acierto_dominio[nombre]:.3f}")

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
for eje, (nombre, X_d) in zip(ejes.flat, dominios.items()):
    eje.imshow(X_d[0, 0], cmap="gray", vmin=0, vmax=1)
    eje.set_title(f"{nombre}: {acierto_dominio[nombre] * 100:.0f} %",
                  fontsize=14)
    eje.axis("off")
plt.show()

promedio_ent = X_ent.mean(dim=0)          # la foto promedio, (1, 28, 28)
for nombre, X_d in dominios.items():
    with torch.no_grad():
        # max devuelve los valores y sus posiciones; .values, los valores
        confianza = red(X_d).softmax(dim=1).max(dim=1).values
    distancia = (X_d.mean(dim=0) - promedio_ent).abs().mean()
    print(f"{nombre:10} confianza media {confianza.mean():.2f}, "
          f"distancia {distancia:.3f}")

import copy

X_rot = torch.rot90(X_ent[:500], 1, dims=(2, 3))
y_rot = y_ent[:500]


def ajustar_copia(X, y, epocas=5):
    copia = copy.deepcopy(red)
    opt = torch.optim.Adam(copia.parameters(), lr=1e-3)
    entrenar(copia, lotes(X, y, mezclar=True), lotes(X_val, y_val),
             epocas, opt, mostrar=False)
    for nombre in ["Normal", "Rotada"]:
        _, exactitud = evaluar(copia, lotes(dominios[nombre], y_val),
                               funcion_perdida)
        print(f"  {nombre}: {exactitud:.3f}")


print("Solo con las 500 rotadas:")
ajustar_copia(X_rot, y_rot)
print("Con 500 rotadas y 500 normales:")
ajustar_copia(torch.cat([X_rot, X_ent[500:1000]]),
              torch.cat([y_rot, y_ent[500:1000]]))

en_A_ent, en_A_val = y_ent < 5, y_val < 5
# ~ invierte la máscara: las fotos de las clases 5 a 9
tarea_A = (X_ent[en_A_ent], y_ent[en_A_ent])
tarea_B = (X_ent[~en_A_ent], y_ent[~en_A_ent])
val_A = lotes(X_val[en_A_val], y_val[en_A_val])
val_B = lotes(X_val[~en_A_val], y_val[~en_A_val])


def por_tareas(modelo):
    _, exactitud_A = evaluar(modelo, val_A, funcion_perdida)
    _, exactitud_B = evaluar(modelo, val_B, funcion_perdida)
    return exactitud_A, exactitud_B


def en_secuencia(datos_B, congelar=False, epocas=3):
    """Tarea A y después datos_B, midiendo las dos tareas por época."""
    red_s = red_nueva()
    opt = torch.optim.Adam(red_s.parameters(), lr=1e-3)
    # *tarea_A pasa las fotos y las etiquetas como dos argumentos
    cargador_A = lotes(*tarea_A, mezclar=True)
    curva = []
    for _ in range(epocas):
        entrenar(red_s, cargador_A, val_A, 1, opt, mostrar=False)
        curva.append(por_tareas(red_s))
    if congelar:
        # La primera capa lineal, red_s[1], deja de cambiar
        for p in red_s[1].parameters():
            p.requires_grad = False
        opt = torch.optim.Adam(red_s[3].parameters(), lr=1e-3)
    cargador_B = lotes(*datos_B, mezclar=True)
    for _ in range(epocas):
        entrenar(red_s, cargador_B, val_B, 1, opt, mostrar=False)
        curva.append(por_tareas(red_s))
    return curva


sin_remedio = en_secuencia(tarea_B)
for epoca, (a, b) in enumerate(sin_remedio, start=1):
    print(f"Época {epoca}: tarea A {a:.3f}, tarea B {b:.3f}")

X_A, y_A = tarea_A
curvas_repaso = {}
for parte in [0.05, 0.10, 0.20]:
    k = int(parte * len(X_A))           # cuántas fotos viejas se guardan
    con_repaso = (torch.cat([tarea_B[0], X_A[:k]]),
                  torch.cat([tarea_B[1], y_A[:k]]))
    curvas_repaso[parte] = en_secuencia(con_repaso)
    a, b = curvas_repaso[parte][-1]
    print(f"Repaso del {parte:.0%} ({k} fotos): "
          f"tarea A {a:.3f}, tarea B {b:.3f}")

a, b = en_secuencia(tarea_B, congelar=True)[-1]
print(f"Congelar la primera capa: tarea A {a:.3f}, tarea B {b:.3f}")

epocas_x = range(1, 7)
plt.figure()
plt.plot(epocas_x, [a for a, b in sin_remedio], "o-", color="C0",
         label="A, sin repaso")
plt.plot(epocas_x, [b for a, b in sin_remedio], "s--", color="C2",
         label="B, sin repaso")
plt.plot(epocas_x, [a for a, b in curvas_repaso[0.20]], "o-",
         color="C1", label="A, repaso 20 %")
# Solo para dibujar: la línea donde termina la tarea A y empieza la B
plt.axvline(3.5, color="C2", linewidth=1)
plt.ylim(-0.05, 1.05)
plt.xlabel("Época")
plt.ylabel("Exactitud en validación")
plt.title("La tarea B empieza en la época 4")
plt.legend(loc="center left", bbox_to_anchor=(0, 0.42), fontsize=11,
           handlelength=1.2)
plt.show()