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()
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()