Modelos de lenguaje: de predecir la siguiente letra a los LLM

import math
import textwrap
import urllib.request
from pathlib import Path

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

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

ruta = Path("datos/quijote-pg2000.txt")
if not ruta.exists():
    ruta.parent.mkdir(exist_ok=True)
    url = "https://www.gutenberg.org/cache/epub/2000/pg2000.txt"
    urllib.request.urlretrieve(url, ruta)

crudo = ruta.read_text(encoding="utf-8").replace("\r\n", "\n")
# El libro está entre estas dos marcas de Project Gutenberg
inicio = crudo.index("*** START OF THE PROJECT GUTENBERG EBOOK")
inicio = crudo.index("\n", inicio) + 1
fin = crudo.index("*** END OF")
texto = crudo[inicio:fin].strip()

print(f"{len(texto):,} caracteres".replace(",", " "))
pos = texto.index("En un lugar de la Mancha")
print(texto[pos:pos + 70])

letras = sorted(set(texto))
V = len(letras)                       # tamaño del vocabulario
a_id = {c: i for i, c in enumerate(letras)}


def codificar(s):
    return [a_id[c] for c in s]


def decodificar(ids):
    return "".join(letras[i] for i in ids)


datos = torch.tensor(codificar(texto))
# El 90 % del libro para entrenar y el último 10 % para validar
corte = int(0.9 * len(datos))
entreno, valid = datos[:corte], datos[corte:]

print("Caracteres distintos:", V)
print("Sancho ->", codificar("Sancho"))
print(f"Azar puro: pérdida ln({V}) = {math.log(V):.2f}")

T = 64        # ventana de contexto: cuántos caracteres mira el modelo
B = 32        # ventanas por lote
gen_lotes = torch.Generator().manual_seed(0)


def lote(fuente, gen=gen_lotes):
    """B ventanas al azar de `fuente` y las mismas corridas un lugar."""
    ix = torch.randint(len(fuente) - T - 1, (B,), generator=gen)
    x = torch.stack([fuente[i:i + T] for i in ix])
    y = torch.stack([fuente[i + 1:i + T + 1] for i in ix])
    return x.to(device), y.to(device)


x, y = lote(entreno)
print(x.shape, y.shape)
print("Entrada: ", repr(decodificar(x[0, :24].tolist())))
print("Objetivo:", repr(decodificar(y[0, :24].tolist())))

class Bloque(nn.Module):
    """Un bloque transformer decodificador (tema 23), en versión corta."""

    def __init__(self, d, cabezas):
        super().__init__()
        self.cabezas = cabezas
        self.norma1 = nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)     # consultas, claves y valores
        self.salida = nn.Linear(d, d)
        self.norma2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(),
                                 nn.Linear(4 * d, d))

    def forward(self, x):
        b, t, d = x.shape
        h = self.cabezas
        q, k, v = self.qkv(self.norma1(x)).split(d, dim=2)
        # (b, t, d) -> (b, cabezas, t, d / cabezas)
        q, k, v = [z.view(b, t, h, d // h).transpose(1, 2)
                   for z in (q, k, v)]
        pesos = q @ k.transpose(-2, -1) / math.sqrt(d // h)
        # Máscara causal: nadie mira a las posiciones de más adelante
        futuro = torch.ones(t, t, dtype=torch.bool, device=x.device)
        pesos = pesos.masked_fill(futuro.triu(1), float("-inf"))
        mezcla = pesos.softmax(dim=-1) @ v
        mezcla = mezcla.transpose(1, 2).reshape(b, t, d)
        x = x + self.salida(mezcla)            # conexión residual
        return x + self.mlp(self.norma2(x))    # conexión residual


class MiniGPT(nn.Module):
    def __init__(self, V, T, d=96, cabezas=4, capas=4):
        super().__init__()
        self.letra = nn.Embedding(V, d)       # un vector por carácter
        self.posicion = nn.Embedding(T, d)    # un vector por posición
        self.bloques = nn.Sequential(
            *[Bloque(d, cabezas) for _ in range(capas)])
        self.norma = nn.LayerNorm(d)
        self.cabeza = nn.Linear(d, V)         # un logit por carácter

    def forward(self, ids):
        pos = torch.arange(ids.shape[1], device=ids.device)
        x = self.letra(ids) + self.posicion(pos)
        return self.cabeza(self.norma(self.bloques(x)))


modelo = MiniGPT(V, T).to(device)
total = sum(p.numel() for p in modelo.parameters())
print(f"Parámetros: {total:,}".replace(",", " "))
logits = modelo(x)
print("Logits:", tuple(logits.shape))
perdida = F.cross_entropy(logits.reshape(-1, V), y.reshape(-1))
print(f"Pérdida antes de entrenar: {perdida.item():.2f}")

@torch.no_grad()
def generar(modelo, comienzo, n=120, temperatura=0.8, semilla=0):
    gen = torch.Generator(device=device).manual_seed(semilla)
    ids = torch.tensor([codificar(comienzo)], device=device)
    for _ in range(n):
        # Solo los últimos T caracteres: la ventana del modelo
        logits = modelo(ids[:, -T:])[:, -1, :]
        if temperatura == 0:
            siguiente = logits.argmax(dim=-1, keepdim=True)
        else:
            probs = F.softmax(logits / temperatura, dim=-1)
            siguiente = torch.multinomial(probs, 1, generator=gen)
        ids = torch.cat([ids, siguiente], dim=1)
    return decodificar(ids[0].tolist())


def mostrar(s):
    """Imprime el texto en renglones de 72 caracteres."""
    print(textwrap.fill(s.replace("\n", " "), 72))


mostrar(generar(modelo, "En un lugar ", n=60))

@torch.no_grad()
def perdida_valid(modelo, lotes=10):
    gen = torch.Generator().manual_seed(1)    # siempre los mismos lotes
    total = 0.0
    for _ in range(lotes):
        xv, yv = lote(valid, gen)
        total += F.cross_entropy(modelo(xv).reshape(-1, V),
                                 yv.reshape(-1)).item()
    return total / lotes


optimizador = torch.optim.AdamW(modelo.parameters(), lr=3e-3)
PASOS = 2000
historia = []             # (paso, pérdida de entrenamiento, de validación)
muestras = {0: generar(modelo, "En un lugar ")}

for paso in range(1, PASOS + 1):
    x, y = lote(entreno)
    perdida = F.cross_entropy(modelo(x).reshape(-1, V), y.reshape(-1))
    optimizador.zero_grad()
    perdida.backward()
    optimizador.step()
    if paso % 250 == 0:
        historia.append((paso, perdida.item(), perdida_valid(modelo)))
    if paso in (500, PASOS):
        muestras[paso] = generar(modelo, "En un lugar ")

for paso, ent, val in historia[::2] + historia[-1:]:
    print(f"paso {paso:4d}   entrenamiento {ent:.3f}"
          f"   validación {val:.3f}")

pasos = [h[0] for h in historia]
plt.plot(pasos, [h[1] for h in historia], color="C2",
         label="Entrenamiento")
plt.plot(pasos, [h[2] for h in historia], color="C0", label="Validación")
plt.scatter(pasos[-1], historia[-1][2], color="C1", zorder=3, s=40)
plt.axhline(math.log(V), color="C2", linestyle="--", linewidth=1)
plt.text(pasos[0], math.log(V) - 0.15, "Azar: ln 92", color="C2",
         va="top")
plt.ylim(0, 5)
plt.xlabel("Paso de entrenamiento")
plt.ylabel("Pérdida por carácter")
plt.title("El mini-GPT aprende el Quijote")
plt.legend()
plt.show()

for paso, muestra in muestras.items():
    print(f"--- paso {paso}")
    mostrar(muestra)

@torch.no_grad()
def siguientes(modelo, comienzo, temperatura=1.0, k=5):
    ids = torch.tensor([codificar(comienzo)], device=device)
    logits = modelo(ids[:, -T:])[0, -1]
    probs = F.softmax(logits / temperatura, dim=-1)
    p, i = probs.topk(k)
    # El espacio se muestra como «␣» para que se vea
    return "  ".join(f"{letras[j].replace(' ', '␣')} {q:.2f}"
                     for j, q in zip(i.tolist(), p.tolist()))


for temp in (0.5, 1.0, 1.5):
    print(f"T = {temp}:  ",
          siguientes(modelo, "En un lugar de la ", temperatura=temp))

mostrar(generar(modelo, "En un lugar ", n=100, temperatura=0))

def recortar(probs, k=None, p=None):
    """Pone en 0 lo que queda fuera del top-k o del top-p y renormaliza."""
    orden, idx = probs.sort(descending=True)
    conservar = torch.ones_like(orden, dtype=torch.bool)
    if k is not None:
        conservar[k:] = False
    if p is not None:
        # Se conserva un token si lo acumulado ANTES de él no llega a p
        antes = orden.cumsum(0) - orden
        conservar &= antes < p
    nuevo = torch.zeros_like(probs)
    nuevo[idx[conservar]] = orden[conservar]
    return nuevo / nuevo.sum()


probs = torch.tensor([0.50, 0.20, 0.15, 0.10, 0.05])
print("top-k = 2:  ", recortar(probs, k=2).numpy().round(3))
print("top-p = 0.8:", recortar(probs, p=0.8).numpy().round(3))

from transformers import AutoModelForCausalLM, AutoTokenizer

nombre = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(nombre)
llm = AutoModelForCausalLM.from_pretrained(nombre, dtype=torch.float32)
llm = llm.to(device).eval()

millones = sum(p.numel() for p in llm.parameters()) / 1e6
print(f"Parámetros: {millones:.0f} millones")
print("Vocabulario:", len(tok), "tokens")
print("Ventana de contexto:", llm.config.max_position_embeddings)
print(tok.tokenize("En un lugar de la Mancha"))

@torch.no_grad()
def logits_siguiente(texto_inicial):
    x = tok(texto_inicial, return_tensors="pt").to(device)
    return llm(**x).logits[0, -1].float().cpu()


inicio_quijote = "En un lugar de la Mancha, de cuyo nombre no quiero"
z = logits_siguiente(inicio_quijote)
favoritos = z.topk(6).indices
for temp in (0.5, 1.0, 1.5):
    p = F.softmax(z / temp, dim=-1)
    # ¿Cuántos tokens hacen falta para juntar el 90 %?
    nucleo = int((p.sort(descending=True).values.cumsum(0) < 0.9).sum()) + 1
    fila = " ".join(f"{tok.decode(i).strip()} {p[i]:.2f}"
                    for i in favoritos)
    print(f"T={temp}: {fila} | top-p: {nucleo}")

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
etiquetas = [tok.decode(i).strip() for i in favoritos]
for eje, temp in zip(ejes.flat, (0.5, 1.0, 1.5)):
    p = F.softmax(z / temp, dim=-1)[favoritos]
    eje.barh(etiquetas[::-1], p.numpy()[::-1], color="C0")
    eje.set_xlim(0, 0.7)
    eje.set_title(f"T = {temp}".replace(".", ","))
# Cuarto panel: tokens que necesita top-p = 0,9 con cada temperatura
nucleos = []
for temp in (0.5, 1.0, 1.5):
    p = F.softmax(z / temp, dim=-1).sort(descending=True).values
    nucleos.append(int((p.cumsum(0) < 0.9).sum()) + 1)
ejes[1, 1].bar(["0,5", "1", "1,5"], nucleos, color="C2")
for i, cuantos in enumerate(nucleos):
    ejes[1, 1].text(i, cuantos * 1.3, cuantos, ha="center")
ejes[1, 1].set_yscale("log")
ejes[1, 1].set_ylim(1, 1e5)
ejes[1, 1].set_title("Top-p 0,9")
ejes[1, 1].set_ylabel("Tokens")
ejes[1, 1].set_xlabel("Temperatura")
fig.suptitle("¿Qué sigue a «…no quiero»?")
plt.show()

mensajes = [{"role": "user", "content": "¿Qué es un modelo de lenguaje?"}]
print(tok.apply_chat_template(mensajes, add_generation_prompt=True,
                              tokenize=False))

@torch.no_grad()
def preguntar(pregunta, sistema=None, max_tokens=60):
    mensajes = [{"role": "user", "content": pregunta}]
    if sistema:
        mensajes.insert(0, {"role": "system", "content": sistema})
    x = tok.apply_chat_template(mensajes, add_generation_prompt=True,
                                return_tensors="pt", return_dict=True)
    x = x.to(device)
    salida = llm.generate(**x, max_new_tokens=max_tokens,
                          do_sample=False)
    # Solo los tokens nuevos: lo que agregó el modelo
    nuevos = salida[0, x["input_ids"].shape[1]:]
    return tok.decode(nuevos, skip_special_tokens=True)


mostrar(preguntar("Explica en una frase qué es un modelo de lenguaje."))

import json

pedido = ("Extrae los datos de la frase y responde solo con JSON con las "
          "claves nombre, ciudad y edad.\n"
          "Frase: Me llamo Ana, tengo 15 años y vivo en Cochabamba.")
respuesta = preguntar(pedido)
# Si el modelo envuelve el JSON en un bloque de Markdown, se lo quitamos
limpio = respuesta.strip().removeprefix("```json").removesuffix("```")
print(limpio.strip())

datos_ana = json.loads(limpio)
print(datos_ana["ciudad"], datos_ana["edad"] + 1)

@torch.no_grad()
def prob_por_token(frase):
    ids = tok(frase, return_tensors="pt").input_ids.to(device)
    logp = llm(ids).logits[0, :-1].float().log_softmax(dim=-1)
    # La probabilidad que le dio a cada token real, mirando los de antes
    reales = ids[0, 1:]
    p = logp.gather(1, reales[:, None]).squeeze(1).exp().cpu()
    return [tok.decode(i) for i in reales], p


for final in ("acordarme.", "cocinar."):
    frase = "En un lugar de la Mancha, de cuyo nombre no quiero " + final
    piezas, p = prob_por_token(frase)
    perplejidad = torch.exp(-p.log().mean()).item()
    print(f"«...no quiero {final}»  perplejidad = {perplejidad:.1f}")

print([(t, round(q, 3)) for t, q in zip(piezas, p.tolist())][-5:])

plt.figure(figsize=(4.8, 4.8))       # más alta: son 16 tokens
etiquetas = [t.strip() or "(espacio)" for t in piezas]
colores = ["C0"] * len(p)
colores[int(p.argmin())] = "C1"          # el token más sorprendente
plt.barh(range(len(p)), (-p.log2()).numpy(), color=colores)
plt.yticks(range(len(p)), etiquetas)
plt.gca().invert_yaxis()
plt.xlabel("Sorpresa: $-\\log_2 p$ (bits)")
plt.title("«...no quiero cocinar.»")
plt.show()

mostrar(preguntar("Resume en dos frases la novela «El cóndor de cristal», "
                  "publicada en 1987.", max_tokens=50))

mostrar(preguntar("¿En qué año se fundó la ciudad de La Paz, en Bolivia? "
                  "Responde en una frase."))
mostrar(preguntar("¿Cuánto es 17 por 23? Responde solo con el número."))
print(17 * 23)