Un GPT desde cero

import warnings
warnings.filterwarnings("ignore")
import math, random, time
from collections import Counter
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

random.seed(0); np.random.seed(0); torch.manual_seed(0)
torch.set_num_threads(max(1, torch.get_num_threads()))

sample_text = (
    "the model learns to predict the next token. the next token depends on the tokens before it. "
    "a token can be a word, a part of a word, or a single character. learning to predict the next "
    "word teaches the model grammar, facts and reasoning, because all of them help prediction. "
    "transformers use attention to let every token look at the tokens that came before it."
)

def bpe_train(text, n_merges):
    seq = list(text)                         # se empieza por los caracteres
    merges = []
    for _ in range(n_merges):
        pairs = Counter(zip(seq, seq[1:]))
        (a, b), count = pairs.most_common(1)[0]
        if count < 2:
            break
        merges.append((a, b))
        new, i = [], 0
        while i < len(seq):                  # se sustituye cada aparición del par
            if i < len(seq) - 1 and seq[i] == a and seq[i + 1] == b:
                new.append(a + b); i += 2
            else:
                new.append(seq[i]); i += 1
        seq = new
    return merges, seq

def bpe_encode(text, merges):
    seq = list(text)
    for a, b in merges:                      # se aplican las fusiones en el orden en que se aprendieron
        new, i = [], 0
        while i < len(seq):
            if i < len(seq) - 1 and seq[i] == a and seq[i + 1] == b:
                new.append(a + b); i += 2
            else:
                new.append(seq[i]); i += 1
        seq = new
    return seq

merges, tokens = bpe_train(sample_text, 40)
print("primeras 15 fusiones:", [a + b for a, b in merges[:15]])
print(f"\n{len(sample_text)} caracteres -> {len(tokens)} tokens tras {len(merges)} fusiones "
      f"({len(sample_text) / len(tokens):.1f} caracteres por token)\n")
print("codificación de una frase nueva:", bpe_encode("the next word predicts the model", merges))

COUNTRIES = {"france": "paris", "spain": "madrid", "italy": "rome", "germany": "berlin", "portugal": "lisbon",
             "japan": "tokyo", "egypt": "cairo", "peru": "lima", "chile": "santiago", "kenya": "nairobi",
             "canada": "ottawa", "norway": "oslo", "greece": "athens", "poland": "warsaw", "cuba": "havana",
             "india": "new delhi", "china": "beijing", "russia": "moscow", "mexico": "mexico city", "ireland": "dublin",
             "austria": "vienna", "sweden": "stockholm", "finland": "helsinki", "turkey": "ankara", "brazil": "brasilia",
             "argentina": "buenos aires", "colombia": "bogota", "costa rica": "san jose", "morocco": "rabat", "australia": "canberra"}

SUBJ = [("i", "yo", 0), ("you", "tú", 1), ("he", "él", 2), ("she", "ella", 2), ("we", "nosotros", 3), ("they", "ellos", 4)]
VERB = {"eat": ("eats", ["como", "comes", "come", "comemos", "comen"]),
        "see": ("sees", ["veo", "ves", "ve", "vemos", "ven"]),
        "want": ("wants", ["quiero", "quieres", "quiere", "queremos", "quieren"]),
        "buy": ("buys", ["compro", "compras", "compra", "compramos", "compran"]),
        "have": ("has", ["tengo", "tienes", "tiene", "tenemos", "tienen"])}
NOUN = [("apple", "manzana", "f"), ("book", "libro", "m"), ("car", "coche", "m"), ("house", "casa", "f"),
        ("dog", "perro", "m"), ("table", "mesa", "f"), ("cat", "gato", "m"), ("shirt", "camisa", "f")]
ADJ = [("red", "rojo", "roja"), ("new", "nuevo", "nueva"), ("old", "viejo", "vieja"), ("white", "blanco", "blanca"), (None, None, None)]

def translation_pairs():
    out = []
    for en_s, es_s, p in SUBJ:
        for v, (v3, conj) in VERB.items():
            for en_n, es_n, g in NOUN:
                for en_a, am, af in ADJ:
                    en = f"{en_s} {v3 if en_s in ('he', 'she') else v} the {en_a + ' ' if en_a else ''}{en_n}"
                    es = f"{es_s} {conj[p]} {'el' if g == 'm' else 'la'} {es_n}" + (f" {am if g == 'm' else af}" if en_a else "")
                    out.append((en, es))
    return out

ALL_TR = translation_pairs()
random.shuffle(ALL_TR)
TR_TEST, TR_TRAIN = ALL_TR[:120], ALL_TR[120:]
ADD_ALL = [(a, b) for a in range(100) for b in range(100)]
random.shuffle(ADD_ALL)
ADD_TEST, ADD_TRAIN = ADD_ALL[:500], ADD_ALL[500:]

def rand_word():
    return "".join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(random.randint(3, 6)))

def example(task, split="train"):
    # devuelve (instrucción, respuesta); una línea de entrenamiento es instrucción + respuesta
    if task == "add":
        a, b = random.choice(ADD_TRAIN if split == "train" else ADD_TEST)
        return f"add {a}+{b}=", f"{a + b}"
    if task == "rev":
        w = rand_word()
        return f"rev {w}=", w[::-1]
    if task == "sort":
        ds = [random.randint(0, 9) for _ in range(random.randint(3, 6))]
        return f"sort {' '.join(map(str, ds))}=", " ".join(map(str, sorted(ds)))
    if task == "es":
        en, es = random.choice(TR_TRAIN if split == "train" else TR_TEST)
        return f"es {en}=", es
    if task == "capital":
        c = random.choice(list(COUNTRIES))
        return f"capital {c}=", COUNTRIES[c]

TASKS = ["add", "rev", "sort", "es", "capital"]
WEIGHTS = [5, 1, 2, 1, 1]          # la aritmética es la habilidad más difícil, así que recibe más datos
corpus = "\n".join("".join(example(random.choices(TASKS, WEIGHTS)[0])) for _ in range(60000)) + "\n"

print(f"corpus: {len(corpus):,} caracteres, {corpus.count(chr(10)):,} líneas\n")
print(corpus[:400])

chars = sorted(set(corpus) | set("".join(e for pair in ALL_TR for e in pair)))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = torch.tensor(encode(corpus), dtype=torch.long)

print(f"vocabulario ({len(chars)} tokens): {''.join(chars)!r}")
print("encode('add 2+2=4') ->", encode("add 2+2=4"))

BLOCK = 64

def get_batch(B, source=None):
    source = data if source is None else source
    ix = torch.randint(len(source) - BLOCK - 1, (B,))
    x = torch.stack([source[i:i + BLOCK] for i in ix])
    y = torch.stack([source[i + 1:i + BLOCK + 1] for i in ix])
    return x, y

x, y = get_batch(1)
for t in range(12):
    print(f"contexto {decode(x[0, :t + 1].tolist())!r:22s} -> objetivo {decode([y[0, t].item()])!r}")

class CausalSelfAttention(nn.Module):
    def __init__(self, d, n_head, block_size):
        super().__init__()
        self.h = n_head
        self.qkv = nn.Linear(d, 3 * d)            # proyecciones Q, K y V en una sola matriz
        self.proj = nn.Linear(d, d)               # W_O
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size, dtype=torch.bool)))
        self.last_attn = None

    def forward(self, x, past=None, keep_attn=False):
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=2)
        q, k, v = (t.view(B, T, self.h, D // self.h).transpose(1, 2) for t in (q, k, v))   # (B, heads, T, dh)
        if past is not None:                       # caché KV: se anteponen las claves/valores de los tokens anteriores
            k, v = torch.cat([past[0], k], dim=2), torch.cat([past[1], v], dim=2)
        present = (k, v)
        if keep_attn:                              # versión explícita, guarda los pesos para dibujarlos
            Tk = k.shape[2]
            att = (q @ k.transpose(-2, -1)) / math.sqrt(D // self.h)
            att = att.masked_fill(~self.mask[Tk - T:Tk, :Tk], float("-inf")).softmax(-1)
            self.last_attn = att.detach()
            y = att @ v
        else:                                      # núcleo fusionado, las mismas matemáticas, más rápido
            y = F.scaled_dot_product_attention(q, k, v, is_causal=(past is None and T > 1))
        return self.proj(y.transpose(1, 2).reshape(B, T, D)), present

class Block(nn.Module):
    def __init__(self, d, n_head, block_size):
        super().__init__()
        self.ln1, self.attn = nn.LayerNorm(d), CausalSelfAttention(d, n_head, block_size)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
    def forward(self, x, past=None, keep_attn=False):
        a, present = self.attn(self.ln1(x), past, keep_attn)
        x = x + a
        return x + self.mlp(self.ln2(x)), present

class GPT(nn.Module):
    def __init__(self, vocab, block_size=64, d=128, n_layer=4, n_head=4):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab, d)
        self.pos_emb = nn.Embedding(block_size, d)
        self.blocks = nn.ModuleList(Block(d, n_head, block_size) for _ in range(n_layer))
        self.ln_f = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.tok_emb.weight                 # pesos atados
        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, (nn.Linear, nn.Embedding)):
            nn.init.normal_(m.weight, 0.0, 0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.zeros_(m.bias)

    def forward(self, idx, targets=None, past=None, keep_attn=False):
        B, T = idx.shape
        start = 0 if past is None else past[0][0].shape[2]
        x = self.tok_emb(idx) + self.pos_emb(torch.arange(start, start + T))
        presents = []
        for i, blk in enumerate(self.blocks):
            x, pr = blk(x, None if past is None else past[i], keep_attn)
            presents.append(pr)
        logits = self.head(self.ln_f(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), targets.view(-1))
        return logits, loss, presents

torch.manual_seed(0)
model = GPT(len(chars), BLOCK, d=128, n_layer=4, n_head=4)
n_params = sum(p.numel() for p in model.parameters())
print(f"GPT con {n_params:,} parámetros (GPT-2 small tiene 124 millones, GPT-3 tiene 175.000 millones)")

@torch.no_grad()
def complete(model, prompt, max_new=40):
    # decodificación voraz hasta un salto de línea: la respuesta a una instrucción de tarea
    idx = torch.tensor([encode("\n" + prompt)])
    out = []
    for _ in range(max_new):
        logits, _, _ = model(idx[:, -BLOCK:])
        nxt = logits[0, -1].argmax().item()
        if itos[nxt] == "\n":
            break
        out.append(nxt)
        idx = torch.cat([idx, torch.tensor([[nxt]])], dim=1)
    return decode(out)

xb, yb = get_batch(64)
with torch.no_grad():
    _, loss0, _ = model(xb, yb)
print(f"pérdida inicial {loss0.item():.3f}   frente a   ln(vocabulario) = {math.log(len(chars)):.3f}")
print("respuesta sin entrenar a 'add 2+2=':", repr(complete(model, "add 2+2=", 20)))

def evaluate(model, n=60, seed=1234):
    state = random.getstate(); random.seed(seed)       # las mismas instrucciones de prueba cada vez
    model.eval()
    res = {}
    for t in TASKS:
        ok = 0
        for _ in range(n):
            p, a = example(t, "test")
            ok += complete(model, p) == a
        res[t] = ok / n
    model.train(); random.setstate(state)
    return res

STEPS, BATCH, EVAL_EVERY = 2500, 32, 500
opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=0.1, betas=(0.9, 0.95))
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=2e-3, total_steps=STEPS, pct_start=0.05)

losses, skill_curve = [], {0: evaluate(model)}
t0 = time.time()
for step in range(1, STEPS + 1):
    xb, yb = get_batch(BATCH)
    _, loss, _ = model(xb, yb)
    opt.zero_grad(); loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step(); sched.step()
    losses.append(loss.item())
    if step % EVAL_EVERY == 0:
        skill_curve[step] = evaluate(model)
        accs = "  ".join(f"{t} {a:.0%}" for t, a in skill_curve[step].items())
        print(f"paso {step:5d}  pérdida {np.mean(losses[-100:]):.3f}  |  {accs}  ({time.time() - t0:.0f} s)", flush=True)

fig, axes = plt.subplots(1, 2, figsize=(15, 4.5))
smooth = np.convolve(losses, np.ones(50) / 50, mode="valid")
axes[0].plot(smooth); axes[0].set_xlabel("paso"); axes[0].set_ylabel("pérdida (media de 50 pasos)")
axes[0].set_title("Pérdida de preentrenamiento: entropía cruzada del siguiente carácter"); axes[0].grid(alpha=0.3)
steps_eval = list(skill_curve)
for t in TASKS:
    axes[1].plot(steps_eval, [skill_curve[s][t] for s in steps_eval], marker="o", label=t)
axes[1].set_xlabel("paso"); axes[1].set_ylabel("exactitud con datos reservados"); axes[1].set_ylim(-0.03, 1.03)
axes[1].set_title("Cada habilidad se aprende a su propio ritmo"); axes[1].legend(); axes[1].grid(alpha=0.3)
plt.tight_layout(); plt.show()

model.eval()
random.seed(42)
for t in TASKS:
    print(f"--- {t} ---")
    for _ in range(3):
        p, a = example(t, "test")
        out = complete(model, p)
        print(f"  {p}{out:28s} {'✓' if out == a else '✗ esperado ' + a}")

for country in ["france", "peru", "japan", "atlantis", "narnia", "germanyland", "nepal", "spainia"]:
    print(f"capital {country:12s} -> {complete(model, f'capital {country}='):14s}"
          f"{'(memorizado)' if country in COUNTRIES else '(nunca visto: alucinado)'}")

probes = [("add 123+4=", "127", "número de 3 cifras (el entrenamiento usó 0..99)"),
          ("rev abcdefghij=", "jihgfedcba", "palabra de 10 letras (el entrenamiento usó 3..6)"),
          ("sort 9 8 7 6 5 4 3 2 1=", "1 2 3 4 5 6 7 8 9", "9 dígitos (el entrenamiento usó 3..6)"),
          ("es the cat eats the red apple=", "?", "sustantivo como sujeto (en el entrenamiento los sujetos eran pronombres)")]
for p, a, why in probes:
    out = complete(model, p)
    print(f"{p:32s} -> {out:22s} esperado {a:20s} [{why}]")

@torch.no_grad()
def next_char_probs(prompt):
    logits, _, _ = model(torch.tensor([encode("\n" + prompt)]))
    return logits[0, -1].softmax(-1)

fig, axes = plt.subplots(1, 2, figsize=(15, 3.5))
for ax, prompt in zip(axes, ["add 23+58=", "\n"]):
    p = next_char_probs(prompt if prompt != "\n" else "")
    top = p.topk(12)
    labels = [repr(itos[i])[1:-1] if itos[i] != " " else "␣" for i in top.indices.tolist()]
    ax.bar(labels, top.values.numpy(), color="tab:blue")
    ax.set_title(f"siguiente carácter tras {prompt!r}" if prompt != "\n" else "primer carácter de una línea nueva (¿qué tarea?)")
    ax.set_ylim(0, 1)
plt.tight_layout(); plt.show()

@torch.no_grad()
def generate(model, prompt="\n", max_new=200, temperature=1.0, top_k=None, top_p=None, seed=0):
    g = torch.Generator().manual_seed(seed)
    idx = torch.tensor([encode(prompt)])
    for _ in range(max_new):
        logits, _, _ = model(idx[:, -BLOCK:])
        logits = logits[0, -1]
        if temperature == 0:
            nxt = logits.argmax().view(1)
        else:
            logits = logits / temperature
            if top_k is not None:
                kth = logits.topk(top_k).values[-1]
                logits = logits.masked_fill(logits < kth, float("-inf"))
            probs = logits.softmax(-1)
            if top_p is not None:
                sorted_p, order = probs.sort(descending=True)
                keep = sorted_p.cumsum(0) - sorted_p < top_p        # el conjunto más pequeño que alcanza top_p
                probs = torch.zeros_like(probs).scatter(0, order[keep], sorted_p[keep])
                probs = probs / probs.sum()
            nxt = torch.multinomial(probs, 1, generator=g)
        idx = torch.cat([idx, nxt.view(1, 1)], dim=1)
    return decode(idx[0].tolist())

for label, kw in [("voraz (T=0)", dict(temperature=0)),
                  ("T=0.5", dict(temperature=0.5)),
                  ("T=1.0, top-p 0.9", dict(temperature=1.0, top_p=0.9)),
                  ("T=2.0 (demasiado caliente)", dict(temperature=2.0))]:
    text = generate(model, "\n", max_new=150, **kw)
    print(f"===== {label} =====\n{text.strip()}\n")

prompt = "add 47+38="
answer = complete(model, prompt)
full = "\n" + prompt + answer
idx = torch.tensor([encode(full)])
with torch.no_grad():
    model(idx, keep_attn=True)
labels = [c if c != "\n" else "⏎" for c in full]
q_rows = list(range(len(full) - len(answer) - 1, len(full) - 1))   # posiciones que predicen cada dígito de la respuesta
print(f"{prompt}{answer}")

L, H = len(model.blocks), model.blocks[0].attn.h
fig, axes = plt.subplots(L, H, figsize=(4.2 * H, 2.5 * L))
for l in range(L):
    for h in range(H):
        A = model.blocks[l].attn.last_attn[0, h][q_rows].numpy()
        ax = axes[l, h]
        ax.imshow(A, cmap="magma", vmin=0, vmax=1, aspect="auto")
        ax.set_xticks(range(len(labels))); ax.set_xticklabels(labels, fontsize=8)
        ax.set_yticks(range(len(q_rows))); ax.set_yticklabels([f"→{answer[i]}" for i in range(len(answer))], fontsize=8)
        ax.set_title(f"capa {l} cabeza {h}", fontsize=9)
fig.suptitle("Atención desde las posiciones que escriben cada dígito de la respuesta (filas) hacia los caracteres anteriores (columnas)")
plt.tight_layout(); plt.show()

@torch.no_grad()
def generate_greedy_nocache(model, prompt, n):
    idx = torch.tensor([encode(prompt)])
    for _ in range(n):
        logits, _, _ = model(idx)                       # se recalcula todo en cada paso
        idx = torch.cat([idx, logits[:, -1].argmax(-1, keepdim=True)], dim=1)
    return decode(idx[0].tolist())

@torch.no_grad()
def generate_greedy_cache(model, prompt, n):
    idx = torch.tensor([encode(prompt)])
    logits, _, past = model(idx)                        # se procesa la instrucción una sola vez
    out = idx
    for _ in range(n):
        nxt = logits[:, -1].argmax(-1, keepdim=True)
        out = torch.cat([out, nxt], dim=1)
        logits, _, past = model(nxt, past=past)         # solo el token nuevo pasa por el modelo
    return decode(out[0].tolist())

prompt, n_new = "\nadd 12+30=42\nrev hello=olleh\n", BLOCK - 32
t0 = time.time(); a = generate_greedy_nocache(model, prompt, n_new); t_no = time.time() - t0
t0 = time.time(); b = generate_greedy_cache(model, prompt, n_new); t_kv = time.time() - t0
print("salida idéntica:", a == b)
print(f"sin caché {t_no * 1000:.0f} ms, con caché KV {t_kv * 1000:.0f} ms ({t_no / t_kv:.1f} veces más rápido)")
with torch.no_grad():
    _, _, past = model(torch.tensor([encode(prompt)]))
print(f"caché: {len(past)} capas x (claves, valores), cada una con forma {tuple(past[0][0].shape)} = (lote, cabezas, tokens hasta ahora, dim. de cabeza)")

def new_task_example(split):
    en, es = random.choice(TR_TRAIN if split == "train" else TR_TEST)
    return f"en {es}=", en

random.seed(7)
ft_corpus = "\n".join("".join(new_task_example("train")) for _ in range(5000)) + "\n"
ft_data = torch.tensor(encode(ft_corpus), dtype=torch.long)
pretrained_state = {k: v.clone() for k, v in model.state_dict().items()}

def finetune(start_from_pretrained, steps=200):
    torch.manual_seed(0)
    m = GPT(len(chars), BLOCK, d=128, n_layer=4, n_head=4)
    if start_from_pretrained:
        m.load_state_dict(pretrained_state)
    o = torch.optim.AdamW(m.parameters(), lr=5e-4 if start_from_pretrained else 2e-3)
    curve = []
    for s in range(steps):
        xb, yb = get_batch(32, ft_data)
        _, loss, _ = m(xb, yb)
        o.zero_grad(); loss.backward(); o.step()
        curve.append(loss.item())
    m.eval()
    random.seed(99)
    acc = np.mean([complete(m, p) == a for p, a in (new_task_example("test") for _ in range(100))])
    return m, curve, acc

t0 = time.time()
ft_model, curve_pre, acc_pre = finetune(True)
_, curve_scratch, acc_scratch = finetune(False)
print(f"ambos ajustados en {time.time() - t0:.0f} s")
print(f"exactitud ES→EN con datos reservados:  desde el preentrenado {acc_pre:.0%}   desde cero {acc_scratch:.0%}")

plt.figure(figsize=(8, 3.8))
plt.plot(np.convolve(curve_pre, np.ones(10) / 10, "valid"), label="partiendo del GPT preentrenado")
plt.plot(np.convolve(curve_scratch, np.ones(10) / 10, "valid"), label="partiendo de pesos aleatorios")
plt.xlabel("paso de ajuste fino"); plt.ylabel("pérdida"); plt.legend(); plt.grid(alpha=0.3)
plt.title("Mismos datos, mismos pasos: el preentrenamiento da ventaja"); plt.show()

random.seed(5)
for _ in range(4):
    p, a = new_task_example("test")
    print(f"  {p}{complete(ft_model, p):24s} esperado: {a}")

after = evaluate(ft_model)
print(f"{'tarea':8s} {'antes del ajuste':>17s} {'después':>8s}")
for t in TASKS:
    print(f"{t:8s} {skill_curve[STEPS][t]:17.0%} {after[t]:8.0%}")

tokens_seen = STEPS * BATCH * BLOCK
ours = 6 * n_params * tokens_seen
rows = [("this notebook", n_params, tokens_seen),
        ("GPT-2 XL", 1.5e9, 1e10),
        ("GPT-3", 175e9, 300e9),
        ("Chinchilla", 70e9, 1.4e12)]
print(f"{'modelo':14s} {'parámetros':>10s} {'tokens':>10s} {'tokens/parám.':>13s} {'FLOPs entren. (6ND)':>20s} {'× el nuestro':>13s}")
for name, N, D in rows:
    print(f"{name:14s} {N:10.2e} {D:10.2e} {D / N:13.1f} {6 * N * D:20.2e} {6 * N * D / ours:13.1e}")