aitana.py

# -*- coding: utf-8 -*-
"""Aitana.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1rBrfs8ReunRZ49FQjyL3Hvx6AJT_qSrs
"""

!pip uninstall -y hf-xet
!pip install -U 'huggingface_hub[cli]'

# Ahora descarga el modelo
!huggingface-cli download gplsi/Aitana-2B-S

import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

model_id = "gplsi/Aitana-2B-S"
tokenizer = AutoTokenizer.from_pretrained(model_id)
vocab_size = len(tokenizer)
generator = pipeline(
    "text-generation",
    model=model_id,
    tokenizer=tokenizer,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
)

# input_text = "Termina la frase en una sola paraula: Açò es or..."
input_text = "El nom correcte de la comunitat autonoma de valencia es..."
generation = generator(
    input_text,
    do_sample=True,
    top_k=10,
    eos_token_id=tokenizer.eos_token_id,
)
print(f"Result: {generation[0]['generated_text']}")

# Método 1: Ver la clase exacta del tokenizer
print("Clase del tokenizer:")
print(type(tokenizer).__name__)
print()

# Método 2: Ver el nombre completo con módulo
print("Clase completa:")
print(type(tokenizer))
print()

# Método 3: Información del modelo base
print("Información del modelo:")
print(f"Model type: {tokenizer.model_max_length}")
print(f"Vocab size: {tokenizer.vocab_size}")
print()

# Método 4: Ver archivos de configuración
print("Nombre del modelo base:")
if hasattr(tokenizer, 'name_or_path'):
    print(tokenizer.name_or_path)
print()

# Método 5: Inspeccionar atributos del tokenizer
print("Atributos especiales:")
if hasattr(tokenizer, 'model_type'):
    print(f"Model type: {tokenizer.model_type}")

# Método 6: Ver toda la configuración
print("\nConfiguración completa:")
print(tokenizer)

colors = [
    '102;194;165', '252;141;98', '141;160;203',
    '231;138;195', '166;216;84', '255;217;47'
]

def show_tokens(sentence: str, model_id: str):
    """ Show the tokens each separated by a different color """

    # Load the tokenizer and tokenize the input
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    token_ids = tokenizer(sentence).input_ids

    # Extract vocabulary length
    print(f"Vocab length: {len(tokenizer)}")

    # Print a colored list of tokens
    for idx, t in enumerate(token_ids):
        print(
            f'\x1b[0;30;48;2;{colors[idx % len(colors)]}m' +
            tokenizer.decode(t) +
            '\x1b[0m', t,
            end=' '
        )

def show_token_ids(sentence: str, model_id: str):
    """ Show the tokens"""

    # Load the tokenizer and tokenize the input
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    token_ids = tokenizer(sentence).input_ids

    # Extract vocabulary length
    print(f"Vocab length: {len(tokenizer)}")

    # Print a colored list of tokens
    for idx, t in enumerate(token_ids):
        print(t)

tokenizer_name = tokenizer.__class__.__name__
print(tokenizer_name)

"""# Arquitectura del modelo"""

# Load the model again to inspect its architecture
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)

# Print the model architecture
print(model)

"""Llama la atención el tamaño del embedding:
Embedding(256000, 2048)

esto implica un vocabulario de 256000 tokens y una dimensión de 2048 (que se mantendrá a lo largo de las transformaciones en el transformer).

Inspeccionamos el método forward para asegurarnos del orden correcto de las capas.
"""

# Imprimir el método forward
import inspect
print(inspect.getsource(model.model.layers[0].__class__.forward))

"""Tenemos (as standard in Llama models):

1-Layer Normalization
2-Self Attention
3-Layer Normalization
4-Multilayer Perceptron (Feed forward layer)

# Inicio del viaje predictivo.
Vamos a analizar por todas las transformaciones que pasa nuestro texto al meterle al modelo la frase:

açò es or...
"""

text = "açò es or..."

show_token_ids(text, model_id)

"""de manera más visual:"""

show_tokens(text, model_id)

"""lo que le pasamos al modelo:"""

token_ids = tokenizer(text).input_ids
print(token_ids)
print(len(token_ids))

"""Ver representación de grafias específicas del valenciano"""

text = "anxova així caparaçò col·legi l'horta mascletà perquè"
show_tokens(text, model_id)

"""Encontramos que "·", "'" y "à" son tokens únicos y que están al final del vocabulario muy cerca del 256000. Vamos a imprimir por curiosidad los 100 últimos tokens"""

for token_id in range(vocab_size - 100, vocab_size):
    token_text = tokenizer.decode([token_id])
    print(f"{token_id}: {repr(token_text)}")

"""por tanto le pasamos al modelo la secuencia
[1, 125647, 594, 785, 889]

La primera capa es el embedding, que (consultando en una tabla de (256000, 2048)nos va a devolver un vector de dimensión 2048 para cada token, es decir un tensor (5, 2048)
"""

tokens = torch.tensor([[1, 125647, 594, 785, 889]])  # Shape: (1, 5)

# Paso 1: Embeddings
with torch.no_grad():
    embeddings = model.model.embed_tokens(tokens)

print("Después de embed_tokens:")
print(f"Shape: {embeddings.shape}")

"""visualizamos los primeros 10 dígitos de cada uno."""

for embedding in embeddings[0][:]:
    print(embedding[:8])

"""de manera más visual:"""

print("Embeddings - Shape: (1, 5, 2048)")
print("=" * 80)

tokens_names = ["<s>", "açò", "es", "or", "..."]

for i in range(5):
    row = embeddings[0][i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"Token {i} ({tokens_names[i]:>4}): [{first_vals}, ..., {last_vals}]  (2048 dims)")

print("=" * 80)

"""Ahora se normaliza a lo largo de las filas. En este caso se hace la RMSNorm, que calcula la raíz cuadrada del promedio de los cuadrados, divide cada valor por él y multiplica por un valor gamma aprendible. Esto es una transformación matemática que no altera la dimensión."""

normalized = model.model.layers[0].input_layernorm(embeddings)
print("\nDespués de input_layernorm (primera capa):")
print(f"Shape: {normalized.shape}")

print("Después de input_layernorm - Shape: (1, 5, 2048)")
print("=" * 80)

tokens_names = ["<s>", "açò", "es", "or", "..."]

for i in range(5):
    row = normalized[0, i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")

print("=" * 80)

"""El siguiente paso es ya la famosa capa de atención. Para saber exactamente qué operaciones se aplican hemos de averiguar qué tipo de atención se usa en esta arquitectura MHA estándar (todos los heads son iguales) o GQA (K y V comparten heads entre grupos de Q)."""

print(f"num_attention_heads: {model.config.num_attention_heads}")
print(f"num_key_value_heads: {model.config.num_key_value_heads}")

"""Tenemos 16 heads. Esto implica que cada head procesa 2048/16 = 128 dimensiones. Por tanto tenemos 16 Queries, 16 Keys y 16 Values (QKV)

La siguiente operación en la transformación lineal Q, donde cada uno de nuestros vectores de 2048 dims se multiplica por una matriz (2048,2048), lo que nos devuelve otro vector de dim 2048. Al tener 5 de ellos no hemos sufrido ninguna transformación en las dimensiones.
"""

Q = model.model.layers[0].self_attn.q_proj(normalized)  # Linear(2048 → 2048)
Q.shape

"""Aquí tenemos la matriz de pesos (adquirida durante el entrenamiento) W_q."""

W_q = model.model.layers[0].self_attn.q_proj.weight

print(f"Shape de los pesos W_q: {W_q.shape}")  # (2048, 2048)
for i in range(5):
    row = W_q[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_q[-5+i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("\n","(2048 d) "*11)

print("=" * 80)

"""Ahora multiplicamos por nuestro tensor con los embeddings normalizados. Lo que nos dará Q con dimensión (5, 2048)"""

all_tokens = normalized[0]  # Shape: (5, 2048)
Q = all_tokens @ W_q.T  # Shape: (5, 2048)
# equivalente a
#Q = model.model.layers[0].self_attn.q_proj(normalized)

print(f"Shape de los pesos W_q: {Q.shape}")  # (2048, 2048)
for i in range(5):
    row = Q[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")

print("=" * 80)

"""De igual forma tenemos la matriz de pesos de las Keys W_k (2048, 2048)"""

W_k = model.model.layers[0].self_attn.k_proj.weight

print(f"Shape de los pesos K_q: {W_k.shape}")  # (2048, 2048)
for i in range(5):
    row = W_k[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_k[-5+i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("\n","(2048 d) "*11)

print("=" * 80)

all_tokens = normalized[0]  # Shape: (5, 2048)
K = all_tokens @ W_k.T  # Shape: (5, 2048)
# equivalente a
#K_all = model.model.layers[0].self_attn.k_proj(normalized)

print(f"Shape de los pesos W_k: {K.shape}")  # (2048, 2048)
for i in range(5):
    row = K[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")

print("=" * 80)

"""y por último y de igual manera los Values, con su W_v y su V"""

W_v = model.model.layers[0].self_attn.v_proj.weight

print(f"Shape de los pesos W_v: {W_v.shape}")  # (2048, 2048)
for i in range(5):
    row = W_v[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_v[-5+i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("\n","(2048 d) "*11)

print("=" * 80)

all_tokens = normalized[0]  # Shape: (5, 2048)
V = all_tokens @ W_v.T  # Shape: (5, 2048)
# equivalente a
#V = model.model.layers[0].self_attn.v_proj(normalized)

print(f"Shape de los pesos W_v: {V.shape}")  # (2048, 2048)
for i in range(5):
    row = V[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")

print("=" * 80)

"""Dividimos los 2048 valores en 16 heads de 128 dimensiones cada uno.

En este modelo cada head tiene su propia Q, K, V completamente independiente. Los 16 heads trabajan en paralelo y luego se concatenan.

Así hay más capacidad en el modelo, 16 proyecciones lineales independientes en paralelo dan más flexibilidad (en otras arquitecturas se tiene un solo Q, K y V, aunque ultimamente se agrupan por grupos como punto medio).

Habrás oído decir que cada head aprende conceptos diferentes como semántica, sintaxis... pero en realidad no se sabe qué aprende cada head ni siquiera si son redundantes unos con otros.

La operació que se realiza es simplemente dividir secuancialmente los 2048 en 16 trozos de 128.
"""

# Configuración
num_heads = 16
head_dim = 2048 // 16  # = 128
print(f'Número de heads: {num_heads}')
print(f'Dimensión de cada head {head_dim}')
# Reshape: dividir 2048 en (16 heads × 128 dims)
Q_reshaped = Q.view(5, num_heads, head_dim)  # (5, 16, 128)
K_reshaped = K.view(5, num_heads, head_dim)  # (5, 16, 128)
V_reshaped = V.view(5, num_heads, head_dim)  # (5, 16, 128)

print(f"\nDespués del reshape:")
print(f"Q: {Q_reshaped.shape}  # (tokens, heads, head_dim)")
print(f"K: {K_reshaped.shape}")
print(f"V: {V_reshaped.shape}")

"""Transponemos de (tokens, heads, dim) a (heads, tokens, dim).

Esto nos permite que cada uno de los 16 heads procese su secuencia completa de forma independiente y en paralelo. Es decir, head 0 procesa sus 5 tokens, head 1 procesa sus 5 tokens, etc.
"""

# Transpose: intercambiar dimensiones (tokens, heads) → (heads, tokens)
Q_heads = Q_reshaped.transpose(0, 1)  # (16, 5, 128)
K_heads = K_reshaped.transpose(0, 1)  # (16, 5, 128)
V_heads = V_reshaped.transpose(0, 1)  # (16, 5, 128)

print(f"\nDespués del transpose:")
print(f"Q: {Q_heads.shape}  # (heads, tokens, head_dim)")
print(f"K: {K_heads.shape}")
print(f"V: {V_heads.shape}")

print(f"16 heads, cada uno con 5 tokens × 128 dimensiones:\n")

tokens_names = ["<s>", "açò", "es", "or", "..."]

for head_idx in range(16):
    print(f"HEAD {head_idx} - Shape: (5, 128)")
    print(f"{'='*80}")

    for token_idx in range(5):
        token_vec = Q_heads[head_idx, token_idx, :]  # (128,)

        first_vals = ", ".join([f"{v:.4f}" for v in token_vec[:5]])
        last_vals = ", ".join([f"{v:.4f}" for v in token_vec[-5:]])

        print(f"Token {token_idx} ({tokens_names[token_idx]:>4}): [{first_vals}, ..., {last_vals}]  (128 dims)")

    print()

"""Transponemos K para poder hacer la multiplicación matricial Q @ K^T.

Queremos comparar cada query (128 dims) con cada key (128 dims) mediante producto escalar. Al transponer K de (16, 5, 128) a (16, 128, 5), la multiplicación Q @ K^T nos da una matriz (16, 5, 5) donde cada posición (i,j) es el score de cuánto el token i atiende al token j.

De esta manera tenemos 16 matrices 5x5
"""

token_vec

# Para cada head, multiplicar Q por K transpuesta
attention_scores = Q_heads @ K_heads.transpose(-2, -1)  # (16, 5, 5)

print(f"Attention scores shape: {attention_scores.shape}")
for head_idx in range(16):
    print(f"HEAD {head_idx}  Scores Matrix - Shape: (5, 5)")
    print(f"{'='*80}")

    for token_idx in range(5):
        token_vec = attention_scores[head_idx, token_idx, :]  # (5,)
        vals = ", ".join([f"{v:.4f}" for v in token_vec[:]])
        print(f"Token {token_idx}({tokens_names[token_idx]:>4}): [{vals}]  (5 dims)")

    print()

"""Aplicamos una máscara a los attention scores para que los tokens no tengan información de los tokens siguientes.

Así cada secuencia sirve para entrenar el modelo a predecir el siguiente token (sin hacer trampa viendo el futuro).
"""

# Crear máscara triangular inferior (True = permitido, False = bloqueado)
causal_mask = torch.tril(torch.ones(5, 5)).bool()

print("Máscara causal (5×5):")
print(causal_mask.int())  # 1 = puede atender, 0 = bloqueado

"""Ponemos `-inf` donde no queremos que el token atienda. Tras el softmax que haremos después, `-inf` se convierte en 0 (peso nulo)."""

# Aplicar máscara: poner -inf donde no puede atender
attention_scores_masked = attention_scores.masked_fill(~causal_mask, float('-inf'))

print(f"Attention scores masked shape: {attention_scores_masked.shape}")
for head_idx in range(16):
    print(f"HEAD {head_idx}  Scores Matrix Masked - Shape: (5, 5)")
    print(f"{'='*80}")

    for token_idx in range(5):
        token_vec = attention_scores_masked[head_idx, token_idx, :]  # (5,)
        vals = ", ".join([f"{v:.4f}" for v in token_vec[:]])
        print(f"Token {token_idx}({tokens_names[token_idx]:>4}): [{vals}]")

    print()

"""Como podemos observar tenemos valores grandes en las matrices (de hasta >100 en head 10).
El siguiente paso es un escalado se divididen los scores por √(head_dim) para estabilizar los gradientes. Podemos observar valores mucho más moderados en todos los heads.
"""

import math

# Escalar los scores
scale = math.sqrt(head_dim)  # sqrt(128) ≈ 11.31
attention_scores_scaled = attention_scores_masked / scale

print(f"Escalado por √{head_dim} = {scale:.2f}")

print(f"Attention scores masked shape: {attention_scores_scaled.shape}")
for head_idx in range(16):
    print(f"HEAD {head_idx}  Scores Matrix Masked - Shape: (5, 5)")
    print(f"{'='*80}")

    for token_idx in range(5):
        token_vec = attention_scores_scaled[head_idx, token_idx, :]  # (5,)
        vals = ", ".join([f"{v:.4f}" for v in token_vec[:]])
        print(f"Token {token_idx}({tokens_names[token_idx]:>4}): [{vals}]")

    print()

"""Aplicamos la función softmax para transformar los scores en distribuciones de probabilidad.

Cada fila representa cuánto atiende ese token a cada posición (suma 1). Los `-inf` se convierten en 0 (no atiende a tokens futuros).

"""

# Aplicar softmax a lo largo de la última dimensión (sobre los keys)
attention_weights = torch.softmax(attention_scores_scaled, dim=-1)
# se puede acceder también haciendo un forward directo con el parámetro output_attentions = True
# outputs = model(
    # input_ids=tokens,
    # output_attentions=True)

# #Acceder a los pesos de atención de la primera capa
# attention_weights_from_model = outputs.attentions[0]

# print(f"Attention weights shape: {attention_weights.shape}  # (16, 5, 5)")
# for i in range(16):
#     print(f"\nHead {i} - Atenciones:")
#     print(attention_weights[i])

print(f"Attention scores masked shape: {attention_weights.shape}")
for head_idx in range(16):
    print(f"HEAD {head_idx}  Scores Matrix Masked - Shape: (5, 5)")
    print(f"{'='*80}")

    for token_idx in range(5):
        token_vec = attention_weights[head_idx, token_idx, :]  # (5,)
        vals = ", ".join([f"{v:.4f}" for v in token_vec[:]])
        print(f"Token {token_idx}({tokens_names[token_idx]:>4}): [{vals}]")

    print()

"""Se puede observar que cada fila suma 1 por lo que es una distribución de probabilidad. Ejecuta esta celda para obtener una fila aleatoria de entre todos los heads."""

import random

head_idx = random.randrange(0,16)
token_idx = random.randrange(0,5)

token_vec = attention_weights[head_idx, token_idx, :]
vals = ", ".join([f"{v:.4f}" for v in token_vec[:]])
print(f"Visualización: fila {token_idx} del Head {head_idx} [{vals}]\n")
print(f"Verificación: fila 2 del Head 0 suma = {attention_weights[head_idx, token_idx].sum():.4f}\n")

"""Multiplicamos ahora por V (Values). Usamos los pesos de atención para hacer una suma ponderada de los valores (V).

Para cada token: su output es la combinación de los valores de todos los tokens a los que puede atender, ponderados por los pesos de atención.

Shape: (16, 5, 5) @ (16, 5, 128) → (16, 5, 128)
"""

attention_output = attention_weights @ V_heads  # (16, 5, 5) @ (16, 5, 128) = (16, 5, 128)

print(f"Attention output shape: {attention_output.shape}")
for head_idx in range(16):
    print(f"HEAD {head_idx}  Scores Matrix Masked - Shape: (5, 5)")
    print(f"{'='*80}")

    for token_idx in range(5):
        token_vec = attention_output[head_idx, token_idx, :]  # (5,)
        vals = ", ".join([f"{v:.4f}" for v in token_vec[:]])
        print(f"Token {token_idx}({tokens_names[token_idx]:>4}): [{vals}]")

    print()

"""Ahora concatenamos todos los heads de vuelta para volver a nuestra matriz (5, 2048). Y además hacemos otra trnasformación lineal (o_proj (2048 → 2048)) para mezclar la información de todos los heads."""

# Paso 1: Transpose - volver a poner tokens primero
attention_output_transposed = attention_output.transpose(0, 1)  # (16, 5, 128) → (5, 16, 128)

print(f"Después del transpose: {attention_output_transposed.shape}")

# Paso 2: Reshape/Concatenar - unir los 16 heads de 128 → 2048
attention_output_concat = attention_output_transposed.reshape(5, 2048)

print(f"Después de concatenar: {attention_output_concat.shape}  # (5, 2048)")


for token_idx in range(5):
    token_vec = attention_output_concat[token_idx, :]  # (2048,)

    first_vals = ", ".join([f"{v:.4f}" for v in token_vec[:5]])
    last_vals = ", ".join([f"{v:.4f}" for v in token_vec[-5:]])

    print(f"Token {token_idx} ({tokens_names[token_idx]:>4}): [{first_vals}, ..., {last_vals}]  (2048 dims)")

"""Pesos de la operación o_proj (proyección de salida)"""

# Mostrar los pesos de o_proj
W_o = model.model.layers[0].self_attn.o_proj.weight

print(f"Matriz de pesos W_o: {W_o.shape}  # (2048, 2048)")
print(f"Shape de los pesos W_o: {W_o.shape}")  # (2048, 2048)
for i in range(5):
    row = W_o[i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_o[-5+i]  # Fila i con 2048 valores

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  (2048 dims)")
print("\n","(2048 d) "*11)

print("=" * 80)

"""Hacemos la última transformación para tener la atención final."""

attention_final = attention_output_concat @ W_o.T

print(f"Después de o_proj: {attention_final.shape}  # (5, 2048)")
for token_idx in range(5):
    token_vec = attention_final[token_idx, :]  # (2048,)

    first_vals = ", ".join([f"{v:.4f}" for v in token_vec[:5]])
    last_vals = ", ".join([f"{v:.4f}" for v in token_vec[-5:]])

    print(f"Token {token_idx} ({tokens_names[token_idx]:>4}): [{first_vals}, ..., {last_vals}]  (2048 dims)")

"""Con esto hemos completado la primera etapa del plan en 5 puntos para el bloque de atención:
Siguiendo el plan, lo siguiente es otr normalización en la capa:\
✓ Atención\
→ Residual connection\
→ post_attention_layernorm (RMSNorm)\
→ MLP\
→ Residual connection

Siguiente paso: Residual connection (sumar con la entrada original antes de normalizar). Esto ayuda al flujo de gradientes durante el entrenamiento (evita vanishing gradients) y permite entrenar redes muy profundas.
"""

# Residual connection
hidden_states = normalized[0] + attention_final

print(f"Después de residual connection: {hidden_states.shape}")
for token_idx in range(5):
    token_vec = hidden_states[token_idx, :]  # (2048,)

    first_vals = ", ".join([f"{v:.4f}" for v in token_vec[:5]])
    last_vals = ", ".join([f"{v:.4f}" for v in token_vec[-5:]])

    print(f"Token {token_idx} ({tokens_names[token_idx]:>4}): [{first_vals}, ..., {last_vals}]  (2048 dims)")

"""Siguiendo el plan, lo siguiente es otr normalización en la capa:\
✓ Atención\
✓ Residual connection\
→ post_attention_layernorm (RMSNorm)\
→ MLP\
→ Residual connection
"""

# post_attention_layernorm
residual = hidden_states  # Guardamos para la próxima residual (celda 67)
hidden_states_normalized = model.model.layers[0].post_attention_layernorm(hidden_states)

print(f"Después de post_attention_layernorm: {hidden_states_normalized.shape}")
for token_idx in range(5):
    token_vec = hidden_states[token_idx, :]  # (2048,)

    first_vals = ", ".join([f"{v:.4f}" for v in token_vec[:5]])
    last_vals = ", ".join([f"{v:.4f}" for v in token_vec[-5:]])

    print(f"Token {token_idx} ({tokens_names[token_idx]:>4}): [{first_vals}, ..., {last_vals}]  (2048 dims)")



"""Vamoms ahora a una capa Multilayer Perceptron(MLP):\
✓ Atención\
✓ Residual connection\
✓ post_attention_layernorm (RMSNorm)\
→ MLP\
→ Residual connection

Esto es una combinación de un aumento de la dimensión y una reducción, en este caso subimos a 5440. La elección de este número es un poco técnica pero la explico para no dejar cabos sueltos. Tiene que ver con optimización de hardware y conocimiento ya adquirido de entrenar otras estructuras.
Este MLP usa la activación SwiGLU, esto hace que tengamos 3 matrices lineales aprendidas (gate_proj, up_proj, down_proj) en lugar de las dos tradicionales.
Se intenta entonces mantener el mismo número de parámetros que se haría con dos capas (que es hacer un x4 en las dims), por tanto:\

MLP estándar (2 capas):  2 × (hidden × 4×hidden) = 8 × hidden²\
SwiGLU (3 capas):        3 × (hidden × intermediate) \


Para igualar FLOPs:\
3 × intermediate ≈ 8 × hidden\
intermediate ≈ (8/3) × hidden ≈ 2.67 × hidden


Para nuestro modelo:
hidden_size = 2048\
8/3 × 2048 ≈ 5461\
Redondeado a múltiplo cercano de potencia de 2: 5440 = 85 × 64

Vamos a mostrar a continuación las tres matrices:\
gate_proj: (2048 → 5440)\
up_proj: (2048 → 5440)\
down_proj: (5440 → 2048)
"""

# Por si no conoces esta activación te dejo la estructura aquí:
# Input: x (2048 dims)
#     ↓
# gate_proj(x) → (5440 dims)  ← rama 1
# up_proj(x)   → (5440 dims)  ← rama 2
#     ↓
# Combinar: Swish(gate_proj(x)) ⊙ up_proj(x)  ← multiplicación elemento a elemento
#     ↓
# down_proj → (2048 dims)
# ```

# **Fórmula:**
# ```
# output = down_proj( Swish(gate_proj(x)) ⊙ up_proj(x) )

"""Para entender gate_proj. Vemos la matriz de pesos W_proj y multiplicamos pot la salida de la normalized layer:"""

W_gate = model.model.layers[0].mlp.gate_proj.weight  # (5440, 2048)


print(f"\nW_gate (gate_proj): {W_gate.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = W_gate[i]
    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_gate.shape[1]} dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_gate[-5+i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_gate.shape[1]} dims)")
print("\n",f"({W_gate.shape[0]} dims) "*11)

"""Multiplicamos `hidden_states_normalized` (la salida del post_attention_layernorm) por W_gate:

"""

gate_output = hidden_states_normalized @ W_gate.T  # (5, 2048) @ (2048, 5440) = (5, 5440)

print(f"\nW_up (up_proj): {gate_output.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = gate_output[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({gate_output.shape[1]} dims)")

"""Lo mismo con up_proj:"""

W_up = model.model.layers[0].mlp.up_proj.weight      # (5440, 2048)

print(f"\nW_up (up_proj): {W_up.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = W_up[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_up.shape[1]} dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_up[-5+i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_up.shape[1]} dims)")
print("\n",f"({W_up.shape[0]} dims) "*11)

"""Multiplicamos `hidden_states_normalized` (la salida del post_attention_layernorm) por W_up:"""

up_output = hidden_states_normalized @ W_up.T  # (5, 2048) @ (2048, 5440) = (5, 5440)

print(f"\nW_up (up_proj): {up_output.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = up_output[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({up_output.shape[1]} dims)")

"""Lo mismo con down_proj:"""

W_down = model.model.layers[0].mlp.down_proj.weight  # (2048, 5440)

print(f"\nW_down (down_proj): {W_down.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = W_down[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_down.shape[1]} dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_down[-5+i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_down.shape[1]} dims)")
print("\n",f"({W_down.shape[0]} dims) "*11)

"""Para combinarlas hemos de hacer la activación SwiGLU:

Dimensiones en cada paso:\

Input: (5, 2048)\
gate_proj: (5, 2048) @ (2048, 5440) → (5, 5440)\
up_proj: (5, 2048) @ (2048, 5440) → (5, 5440)\
SwiGLU: Swish + ⊙: (5, 5440) (no cambia dimensión)\
down_proj: (5, 5440) @ (5440, 2048) → (5, 2048)\

Recordamos que Swish es: Swish(x) = x * sigmoid(x)

"""

# hidden_states (5, 2048)
#     ↓
# ├→ × W_gate^T → gate_output (5, 5440)
# │                    ↓
# │               Swish(gate_output) (5, 5440)
# │                    ↓
# └→ × W_up^T   → up_output (5, 5440)
#                      ↓
#         Swish(gate_output) ⊙ up_output  (5, 5440)
#         (multiplicación elemento a elemento)
#                      ↓
#               × W_down^T → output (5, 2048)

# Opción 1: Manual
swish_gate = gate_output * torch.sigmoid(gate_output)

# Opción 2: Usando SiLU (que es exactamente Swish)
# swish_gate = torch.nn.functional.silu(gate_output)

print(f"swish_gate shape: {swish_gate.shape}  # (5, 5440)")
print(f"\n{'='*80}")
for i in range(5):
    row = swish_gate[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({swish_gate.shape[1]} dims)")

"""Siguiente paso: Multiplicar swish_gate ⊙ up_output (elemento a elemento)."""

mlp_intermediate = swish_gate * up_output  # (5, 5440)

print(f"swish_gate shape: {mlp_intermediate.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = mlp_intermediate[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({mlp_intermediate.shape[1]} dims)")

"""Siguiente paso: Multiplicar por down_proj (transpuesta) para volver a 2048 dimensiones."""

# Proyección final del MLP
mlp_output = mlp_intermediate @ W_down.T  # (5, 5440) @ (5440, 2048) = (5, 2048)

print(f"mlp_output shape: {mlp_output.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = mlp_output[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({mlp_output.shape[1]} dims)")

"""Vamoms ahora a una capa Multilayer Perceptron(MLP):\
✓ Atención\
✓ Residual connection\
✓ post_attention_layernorm (RMSNorm)\
✓ MLP\
→ Residual connection

Añadimos ahora otra residual connection, y así damos tick al último elemento de la lista:
"""

# Sumar con el residual guardado antes del MLP
final_output = residual + mlp_output  # (5, 2048)

print(f"Salida final del decoder layer 0: {final_output.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = final_output[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({final_output.shape[1]} dims)")

"""Vale por fin hemos terminado con 1 bloque de attention ¡pero hay 23 más!...
La salida de cada capa se convierte en la entrada de la siguiente. Cada capa refina y transforma la representación de los tokens.

Las 24 capas procesan secuencialmente, manteniendo siempre la dimensión (5, 2048).
"""

hidden.shape

print("="*80)
print("APILAMIENTO DE LAS 24 CAPAS")
print("="*80)

# Usar el modelo completo para obtener las salidas de todas las capas
with torch.no_grad():
    outputs = model(
        input_ids=tokens,
        output_hidden_states=True,  # Esto nos da la salida de cada capa
        return_dict=True
    )

# outputs.hidden_states contiene:
# [0] = embeddings
# [1] = salida capa 0
# [2] = salida capa 1
# ...
# [24] = salida capa 23

print(f"Número de hidden states (embeddings + 24 capas): {len(outputs.hidden_states)}")

# Mostrar evolución de cada capa
for layer_idx in range(25):  # 0=embeddings, 1-24=capas
    hidden = outputs.hidden_states[layer_idx]

    if layer_idx == 0:
        print(f"\nEmbeddings iniciales:")
    else:
        print(f"\nCapa {layer_idx-1:2d} completada:")

    for i in range(5):
        row = hidden[0,i]

        # Primeros 5 valores
        first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

        # Últimos 5 valores
        last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

        print(f"[{first_vals}, ..., {last_vals}]  ({hidden.shape[2]} dims)")

print(f"\n{'='*80}")
print(f"Shape final: {outputs.hidden_states[-1].shape}")

"""Una vez salimos de las 24 capas de atención nos quedan solamente unas pocas capas para llegar al final:
-RMSNorm final: Normaliza la salida de la última capa decoder.
-lm_head: Proyecta de 2048 dimensiones al vocabulario completo (256000). Cada posición ahora tiene un score para cada token posible.
"""

last_layer_output = outputs.hidden_states[-1]
# Paso 1: Final RMSNorm
with torch.no_grad():
    normalized_output = model.model.norm(last_layer_output)

print(f"\nDespués de RMSNorm final: {normalized_output.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = normalized_output[0,i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({normalized_output.shape[2]} dims)")

"""Pasamos ahora a 256000 dims para obtener los logits (valores asignados que corresponden 1 a 1 con los tokens del vocabulario). para ello tenemos la gigantesca matriz de pesos W_lm_head (256000, 2048)"""

W_lm_head = model.lm_head.weight  # (256000, 2048)

print("="*80)
print("MATRIZ LM_HEAD")
print("="*80)

print(f"\nW_lm_head shape: {W_lm_head.shape} ")
print(f"\n{'='*80}")
for i in range(5):
    row = W_lm_head[i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_lm_head.shape[1]} dims)")
print("[", "...     "*12, "]")
for i in range(5):
    row = W_lm_head[-5+i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({W_lm_head.shape[1]} dims)")
print("\n",f"({W_lm_head.shape[0]} d) "*11)

logits = normalized_output @ W_lm_head.T  # (1, 5, 2048) @ (2048, 256000) = (1, 5, 256000)

# # alternativa: lm_head - proyección al vocabulario
# with torch.no_grad():
#     logits = model.lm_head(normalized_output)

print(f"\nDespués de lm_head: {logits.shape}")
print(f"\n{'='*80}")
for i in range(5):
    row = logits[0,i]

    # Primeros 5 valores
    first_vals = ", ".join([f"{v:.4f}" for v in row[:5]])

    # Últimos 5 valores
    last_vals = ", ".join([f"{v:.4f}" for v in row[-5:]])

    print(f"[{first_vals}, ..., {last_vals}]  ({logits.shape[2]} dims)")

"""Estamos muy cerquita ya, para predecir el siguiente token usamos solamente la última fila de esta matriz que corresponde a la predicción para el último token de la secuencia ('...')"""

print("="*80)
print("PREDICCIÓN DEL SIGUIENTE TOKEN")
print("="*80)

# Para predecir el siguiente token, usamos el ÚLTIMO token de la secuencia (posición 4)
last_token_logits = logits[0, -1, :]  # (256000,)

print(f"\nLogits del último token: {last_token_logits.shape}")

# Primeros 5 valores
first_vals = ", ".join([f"{v:.4f}" for v in last_token_logits[:5]])

# Últimos 5 valores
last_vals = ", ".join([f"{v:.4f}" for v in last_token_logits[-5:]])

print(f"[{first_vals}, ..., {last_vals}] ({last_token_logits.shape[0]} dims)")

"""Con esto ya sacamos las probabilidades con una softmax.
Hemos de decir que el modelo llega hasta aquí, al usar el modelo este nos devuelve los logits. Con eso nosotros ya aplicamos softmax y la estrategia de sampleado (elección del token) que querramos. Se usan los siguiente parámetros normalmente:
"""

# output_tokens = model.generate(
#     tokens,
#     max_length=50,
#     temperature=1.0,      # ← AQUÍ se especifica
#     top_p=0.9,           # ← AQUÍ se especifica
#     top_k=50,            # ← AQUÍ se especifica
#     do_sample=True       # ← True = sampling, False = greedy
# )

"""Explico estos valores al final del notebook como apéndice."""

# Aplicar softmax para obtener probabilidades
probabilities = torch.softmax(last_token_logits, dim=-1)

print(f"\nProbabilidades: {probabilities.shape}")
print(f"Suma de probabilidades: {probabilities.sum():.6f}  # Debe ser 1.0")
# Primeros 5 valores
first_vals = ", ".join([f"{v:.4f}" for v in probabilities[:5]])

# Últimos 5 valores
last_vals = ", ".join([f"{v:.4f}" for v in probabilities[-5:]])

print(f"[{first_vals}, ..., {last_vals}] ({probabilities.shape[0]} dims)")

# Obtener el token con mayor probabilidad (greedy decoding)
predicted_token_id = torch.argmax(probabilities).item()
predicted_probability = probabilities[predicted_token_id].item()

print(f"\n{'='*80}")
print(f"TOKEN PREDICHO:")
print(f"  ID: {predicted_token_id}")
print(f"  Texto: {tokenizer.decode([predicted_token_id])}")
print(f"  Probabilidad: {predicted_probability:.4f} ({predicted_probability*100:.2f}%)")

# Top-5 predicciones
top5_probs, top5_ids = torch.topk(probabilities, 5)
print(f"\nTop 5 predicciones:")
for i in range(5):
    token_id = top5_ids[i].item()
    prob = top5_probs[i].item()
    text = tokenizer.decode([token_id])
    print(f"  {i+1}. [{token_id}] '{text}' - {prob:.4f} ({prob*100:.2f}%)")

"""¡HEMOS COMPLETADO EL RECORRIDO COMPLETO! 🎉"""

# Input tokens
#     ↓
# Embeddings (256000 → 2048)
#     ↓
# ┌─ CAPA 0 (y se repite 24 veces) ────┐
# │  input_layernorm (RMSNorm)         │
# │  Multi-Head Attention              │
# │  Residual connection               │
# │  post_attention_layernorm (RMSNorm)│
# │  MLP con SwiGLU                    │
# │  Residual connection               │
# └────────────────────────────────────┘
#     ↓
# ... (capas 1-23 igual) ...
#     ↓
# RMSNorm FINAL (después de las 24 capas)
#     ↓
# lm_head (2048 → 256000)
#     ↓
# Softmax
#     ↓
# Predicción

"""Antes de irte recordamos que esto ha sido solo para predecir el siguiente de ..., pero para predecir el siguiente hay que volver a pasar por todo el modelo incluyendo el nuevo token predicho."""

# Input: [<s>, açò, es, or, ...]
#        ↓ (TODO el modelo)
# Output: predice "es" como siguiente token

# Para seguir generando:
# Paso 1: Input [<s>, açò, es, or, ...] → Predice token es
# Paso 2: Input [<s>, açò, es, or, ..., es] → Volver a pasar por TODO el modelo → Predice siguiente token
# Paso 3: Input [<s>, açò, es, or, ..., es, nuevo_token] → Volver a pasar por TODO → Predice siguiente
# Y así sucesivamente...

"""Cada vez:

Se añade el token predicho al final de la secuencia\
Se pasa toda la secuencia por TODO el modelo de nuevo

  Embeddings\
  24 capas completas (attention + MLP)\
  RMSNorm final\
  lm_head\
  Softmax


Se toma el último logit para predecir el siguiente

Por eso la generación de texto es costosa computacionalmente: cada token nuevo requiere un forward pass completo por las 24 capas.
(Aquí es donde técnicas como KV-cache ayudan a optimizar, pero eso es otro tema)

## APÉNDICE

Explicación parámetros de model.generate():

do_sample (True/False):
"""

# do_sample=False  # Greedy decoding (siempre el token con mayor probabilidad)
# do_sample=True   # Sampling probabilístico (introduce aleatoriedad)

"""temperature (float, 0.1-2.0):"""

# logits_ajustados = logits / temperature
# probabilities = softmax(logits_ajustados)

# temperature=0.5   # Más conservador, menos aleatorio
# temperature=1.0   # Distribución original
# temperature=1.5   # Más creativo, más aleatorio

"""top_k (int):"""

top_k=50  # Solo considera los 50 tokens más probables

# Ordena tokens por probabilidad
# Toma solo los top-k
# Pone probabilidad 0 al resto
# Renormaliza y samplea de esos k

# Efecto: Evita tokens muy improbables/raros

"""top_p (float entre 0 y 1, típicamente 0.9-0.95):"""

top_p=0.9  # Considera el conjunto mínimo de tokens que suman 90% de probabilidad

# Ordena tokens por probabilidad descendente
# Suma probabilidades acumuladas hasta llegar a p
# Solo samplea de ese conjunto "núcleo"

# Ejemplo:

# Token A: 40%, Token B: 30%, Token C: 15%, Token D: 10%, otros: 5%
# Con top_p=0.9 → considera A, B, C, D (suman 95% > 90%)
# Con top_p=0.7 → considera solo A, B (suman 70%)

# Ventaja sobre top_k: Es dinámico. A veces incluye 10 tokens, a veces 100, dependiendo de qué tan concentrada esté la distribución.

"""y por último aunque no tiene que ver con el sampleado tenemos

max_length o max_new_tokens, eos_token_id y min_length para decidir cuándo acabar.
"""

max_length=100        # Máxima longitud TOTAL (input + generado)
max_new_tokens=50     # Máximo número de tokens NUEVOS a generar
eos_token_id:         # para cuando encuentra el token de fin (ej: </s>)
min_length=25         # longitud mínima antes de poder parar