Diseño y aplicación de LSTM

pip install torch torchvision torchtext matplotlib numpy pandas

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, random_split
import numpy as np
import matplotlib.pyplot as plt
import time
import math
import warnings
warnings.filterwarnings('ignore')

SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
print(f"PyTorch: {torch.__version__}")

# Crear una LSTM
lstm = nn.LSTM(
    input_size=10,      # Dimensión de cada vector de entrada x_t
    hidden_size=64,     # Dimensión de h_t y c_t
    num_layers=2,       # Capas LSTM apiladas (stacked)
    batch_first=True,   # Input shape: (batch, seq_len, features) — ¡IMPORTANTE!
    dropout=0.2,        # Dropout entre capas (solo si num_layers > 1)
    bidirectional=False, # True → procesa la secuencia en ambas direcciones
)

# Forward pass
batch_size, seq_len, input_size = 8, 20, 10
x = torch.randn(batch_size, seq_len, input_size)

# output: predictions en cada timestep → (batch, seq_len, hidden_size * num_directions)
# (h_n, c_n): hidden state y cell state finales → (num_layers * num_directions, batch, hidden_size)
output, (h_n, c_n) = lstm(x)

print(f"Input:      {x.shape}")           # [8, 20, 10]
print(f"Output:     {output.shape}")       # [8, 20, 64]
print(f"h_n:        {h_n.shape}")          # [2, 8, 64]  (2 layers)
print(f"c_n:        {c_n.shape}")          # [2, 8, 64]

def train_epoch(model, loader, criterion, optimizer, clip=1.0):
    """Entrena un epoch. Retorna loss media."""
    model.train()
    total_loss, n_batches = 0, 0
    for batch in loader:
        optimizer.zero_grad()
        # Cada loader devuelve un formato distinto — se gestiona en cada ejemplo
        loss = model.compute_loss(batch, criterion)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        total_loss += loss.item()
        n_batches += 1
    return total_loss / n_batches


def evaluate(model, loader, criterion):
    """Evalúa en validación/test. Retorna loss media."""
    model.eval()
    total_loss, n_batches = 0, 0
    with torch.no_grad():
        for batch in loader:
            loss = model.compute_loss(batch, criterion)
            total_loss += loss.item()
            n_batches += 1
    return total_loss / n_batches


def plot_losses(train_losses, val_losses, title="Curvas de entrenamiento"):
    """Dibuja las curvas de loss."""
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(train_losses, label='Train', color='#E17055', linewidth=2)
    ax.plot(val_losses, label='Val', color='#74b9ff', linewidth=2)
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.set_title(title)
    ax.legend()
    ax.grid(alpha=0.15)
    plt.tight_layout()
    plt.show()


def count_params(model):
    """Cuenta parámetros entrenables."""
    n = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Parámetros entrenables: {n:,}")
    return n

# Dataset clásico de airline passengers (Box & Jenkins, 1976)
# 144 observaciones mensuales: enero 1949 - diciembre 1960
raw_data = [
    112,118,132,129,121,135,148,148,136,119,104,118,
    115,126,141,135,125,149,170,170,158,133,114,140,
    145,150,178,163,172,178,199,199,184,162,146,166,
    171,180,193,181,183,218,230,242,209,191,172,194,
    196,196,236,235,229,243,264,272,237,211,180,201,
    204,188,235,227,234,264,302,293,259,229,203,229,
    242,233,267,269,270,315,364,347,312,274,237,278,
    284,277,317,313,318,374,413,405,355,306,271,306,
    315,301,356,348,355,422,465,467,404,347,305,336,
    340,318,362,348,363,435,491,505,404,359,310,337,
    360,342,406,396,420,472,548,559,463,407,362,405,
    417,391,419,461,472,535,622,606,508,461,390,432,
]

data = np.array(raw_data, dtype=np.float32)

# Normalizar a [0, 1]
data_min, data_max = data.min(), data.max()
data_norm = (data - data_min) / (data_max - data_min)

# Visualizar
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(data, color='#E17055', linewidth=1.5)
axes[0].set_title('Airline Passengers (original)')
axes[0].set_xlabel('Mes'); axes[0].set_ylabel('Pasajeros (miles)')
axes[1].plot(data_norm, color='#74b9ff', linewidth=1.5)
axes[1].set_title('Normalizado [0, 1]')
axes[1].set_xlabel('Mes')
plt.tight_layout(); plt.show()

print(f"Samples: {len(data)} | Min: {data.min():.0f} | Max: {data.max():.0f}")

class TimeSeriesDataset(Dataset):
    """Dataset de ventanas deslizantes para series temporales."""
    def __init__(self, series, window_size=24):
        self.x, self.y = [], []
        for i in range(len(series) - window_size):
            self.x.append(series[i:i+window_size])
            self.y.append(series[i+1:i+window_size+1])  # shifted 1 step
        self.x = torch.tensor(np.array(self.x), dtype=torch.float32).unsqueeze(-1)  # (N, W, 1)
        self.y = torch.tensor(np.array(self.y), dtype=torch.float32).unsqueeze(-1)  # (N, W, 1)

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

WINDOW = 24  # 2 años de datos mensuales

# Split: primeros 120 meses = train, últimos 24 = test
train_data = data_norm[:120]
test_data = data_norm[120:]

train_ds = TimeSeriesDataset(train_data, window_size=WINDOW)
train_loader = DataLoader(train_ds, batch_size=16, shuffle=True)

print(f"Train windows: {len(train_ds)}")
print(f"Input shape:   {train_ds[0][0].shape}")  # (24, 1)
print(f"Target shape:  {train_ds[0][1].shape}")   # (24, 1)

class LSTMOneToOne(nn.Module):
    """
    LSTM para predicción autoregresiva one-to-one.
    En entrenamiento: procesa ventana completa (teacher forcing).
    En inferencia: un paso a la vez, realimentando la predicción.
    """
    def __init__(self, input_size=1, hidden_size=64, num_layers=2, dropout=0.1):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
        )
        self.fc = nn.Linear(hidden_size, 1)  # 1 output por timestep

    def forward(self, x, hidden=None):
        """
        x: (batch, seq_len, 1)
        hidden: tuple (h, c) o None
        Returns: predictions (batch, seq_len, 1), (h_n, c_n)
        """
        out, (h_n, c_n) = self.lstm(x, hidden)  # out: (B, T, H)
        pred = self.fc(out)                       # pred: (B, T, 1)
        return pred, (h_n, c_n)

    def compute_loss(self, batch, criterion):
        """Para usar con train_epoch/evaluate genéricos."""
        x, y = batch
        x, y = x.to(device), y.to(device)
        pred, _ = self.forward(x)
        return criterion(pred, y)

model_oto = LSTMOneToOne(input_size=1, hidden_size=64, num_layers=2).to(device)
count_params(model_oto)

criterion = nn.MSELoss()
optimizer = optim.Adam(model_oto.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=10, factor=0.5)

EPOCHS = 100
train_losses = []

for epoch in range(EPOCHS):
    loss = train_epoch(model_oto, train_loader, criterion, optimizer)
    train_losses.append(loss)
    scheduler.step(loss)

    if (epoch + 1) % 25 == 0:
        print(f"Epoch {epoch+1:3d} | Train Loss: {loss:.6f} | LR: {optimizer.param_groups[0]['lr']:.1e}")

plt.figure(figsize=(8, 3))
plt.plot(train_losses, color='#E17055', linewidth=1.5)
plt.xlabel('Epoch'); plt.ylabel('MSE Loss'); plt.title('One-to-One: Entrenamiento')
plt.grid(alpha=0.15); plt.tight_layout(); plt.show()

def autoregressive_predict(model, seed_sequence, n_future, device):
    """
    Predicción one-to-one autoregresiva.
    1. Procesa la seed_sequence para inicializar el hidden state.
    2. Predice n_future pasos, realimentando cada predicción.
    """
    model.eval()
    with torch.no_grad():
        # Fase 1: warm-up con la secuencia semilla
        x = torch.tensor(seed_sequence, dtype=torch.float32).unsqueeze(0).unsqueeze(-1).to(device)
        _, hidden = model(x)  # Solo nos interesa el hidden state final

        # Fase 2: predicción autoregresiva one-to-one
        predictions = []
        current_input = x[:, -1:, :]  # Último valor de la semilla: (1, 1, 1)

        for _ in range(n_future):
            pred, hidden = model(current_input, hidden)  # (1, 1, 1), (h, c)
            predictions.append(pred.item())
            current_input = pred  # Realimentar: el output se convierte en input

    return predictions

# Usar los primeros 120 meses como semilla, predecir los 24 restantes
seed = data_norm[:120]
n_predict = 24

preds_norm = autoregressive_predict(model_oto, seed, n_predict, device)

# Desnormalizar
preds = np.array(preds_norm) * (data_max - data_min) + data_min
actual = data[120:]

# Visualizar
fig, ax = plt.subplots(figsize=(10, 5))
months = np.arange(len(data))
ax.plot(months[:120], data[:120], color='#636e72', linewidth=1, alpha=0.5, label='Train')
ax.plot(months[120:], actual, color='#74b9ff', linewidth=2, label='Real')
ax.plot(months[120:], preds, color='#E17055', linewidth=2, linestyle='--', label='Predicción')
ax.axvline(120, color='rgba(255,255,255,.2)', linestyle=':')
ax.set_xlabel('Mes'); ax.set_ylabel('Pasajeros (miles)')
ax.set_title('One-to-One: Predicción autoregresiva de Airline Passengers')
ax.legend(); ax.grid(alpha=0.1); plt.tight_layout(); plt.show()

# Métricas
mae = np.mean(np.abs(preds - actual))
rmse = np.sqrt(np.mean((preds - actual) ** 2))
print(f"MAE:  {mae:.1f} pasajeros")
print(f"RMSE: {rmse:.1f} pasajeros")

from torchtext.datasets import IMDB
from torchtext.data.utils import get_tokenizer
from collections import Counter

# Tokenizador básico
tokenizer = get_tokenizer('basic_english')

# Construir vocabulario a partir del train set
counter = Counter()
train_iter = IMDB(split='train')
for label, text in train_iter:
    counter.update(tokenizer(text))

# Vocabulario: las 20,000 palabras más frecuentes
VOCAB_SIZE = 20_000
vocab_list = ['<pad>', '<unk>'] + [w for w, _ in counter.most_common(VOCAB_SIZE - 2)]
word2idx = {w: i for i, w in enumerate(vocab_list)}
PAD_IDX = word2idx['<pad>']
UNK_IDX = word2idx['<unk>']

def encode_text(text, max_len=256):
    """Tokeniza y codifica un texto, truncando/padding a max_len."""
    tokens = tokenizer(text)[:max_len]
    indices = [word2idx.get(t, UNK_IDX) for t in tokens]
    # Padding
    indices += [PAD_IDX] * (max_len - len(indices))
    return indices

print(f"Vocabulario: {len(word2idx):,} palabras")
print(f"Ejemplo: {encode_text('this movie is great', 8)}")

MAX_LEN = 256

class IMDBDataset(Dataset):
    def __init__(self, split='train'):
        self.data = []
        for label, text in IMDB(split=split):
            encoded = encode_text(text, MAX_LEN)
            # IMDB labels: 'pos'=1(→1), 'neg'=2(→0) en torchtext
            y = 1.0 if label == 2 else 0.0  # pos=1, neg=0
            self.data.append((torch.tensor(encoded, dtype=torch.long),
                              torch.tensor(y, dtype=torch.float32)))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

# Datasets
train_full = IMDBDataset('train')
test_ds = IMDBDataset('test')

# Split train/val (22,500 / 2,500)
train_ds, val_ds = random_split(train_full, [22500, 2500],
                                generator=torch.Generator().manual_seed(SEED))

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader   = DataLoader(val_ds, batch_size=128, shuffle=False)
test_loader  = DataLoader(test_ds, batch_size=128, shuffle=False)

print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | Test: {len(test_ds)}")

class LSTMManyToOne(nn.Module):
    """
    Many-to-One: Embedding → LSTM → último hidden → FC → sigmoid.
    Para clasificación binaria de secuencias (sentimiento).
    """
    def __init__(self, vocab_size, embed_dim=128, hidden_size=128,
                 num_layers=2, dropout=0.3, pad_idx=0, bidirectional=True):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=bidirectional,
        )
        self.dropout = nn.Dropout(dropout)
        # Si bidirectional: hidden_size * 2
        fc_input = hidden_size * 2 if bidirectional else hidden_size
        self.fc = nn.Linear(fc_input, 1)

    def forward(self, x):
        """
        x: (batch, seq_len) — índices de tokens
        Returns: logits (batch, 1)
        """
        emb = self.dropout(self.embedding(x))      # (B, T, E)
        output, (h_n, c_n) = self.lstm(emb)         # output: (B, T, H*2)

        # ── CLAVE: extraer el hidden state final ──
        # Opción 1: último timestep del output
        # last = output[:, -1, :]

        # Opción 2 (mejor con bidireccional): concatenar h_n de ambas direcciones
        if self.lstm.bidirectional:
            # h_n shape: (num_layers*2, B, H)
            # Últimas 2 entradas = última capa forward + backward
            h_forward = h_n[-2]   # (B, H)
            h_backward = h_n[-1]  # (B, H)
            hidden = torch.cat([h_forward, h_backward], dim=1)  # (B, H*2)
        else:
            hidden = h_n[-1]  # (B, H)

        hidden = self.dropout(hidden)
        logits = self.fc(hidden)  # (B, 1)
        return logits

    def compute_loss(self, batch, criterion):
        x, y = batch
        x, y = x.to(device), y.to(device)
        logits = self.forward(x).squeeze(-1)  # (B,)
        return criterion(logits, y)

model_mto = LSTMManyToOne(
    vocab_size=VOCAB_SIZE, embed_dim=128, hidden_size=128,
    num_layers=2, bidirectional=True, pad_idx=PAD_IDX
).to(device)
count_params(model_mto)

criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model_mto.parameters(), lr=1e-3)

EPOCHS = 8
train_losses, val_losses = [], []

for epoch in range(EPOCHS):
    t_loss = train_epoch(model_mto, train_loader, criterion, optimizer)
    v_loss = evaluate(model_mto, val_loader, criterion)
    train_losses.append(t_loss)
    val_losses.append(v_loss)
    print(f"Epoch {epoch+1} | Train: {t_loss:.4f} | Val: {v_loss:.4f}")

plot_losses(train_losses, val_losses, "Many-to-One: Sentimiento IMDB")

# ── Accuracy en test ──
model_mto.eval()
correct, total = 0, 0
with torch.no_grad():
    for x, y in test_loader:
        x, y = x.to(device), y.to(device)
        logits = model_mto(x).squeeze(-1)
        preds = (torch.sigmoid(logits) >= 0.5).float()
        correct += (preds == y).sum().item()
        total += y.size(0)

print(f"\n🎯 Test Accuracy: {correct/total:.2%}")

import urllib.request
import os

# Descargar Tiny Shakespeare (~1MB)
DATA_PATH = 'shakespeare.txt'
if not os.path.exists(DATA_PATH):
    url = 'https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt'
    urllib.request.urlretrieve(url, DATA_PATH)

with open(DATA_PATH, 'r', encoding='utf-8') as f:
    text = f.read()

print(f"Longitud total: {len(text):,} caracteres")
print(f"Primeros 200 chars:\n{text[:200]}")

# Vocabulario de caracteres
chars = sorted(set(text))
CHAR_VOCAB_SIZE = len(chars)
char2idx = {c: i for i, c in enumerate(chars)}
idx2char = {i: c for c, i in char2idx.items()}

print(f"\nVocabulario: {CHAR_VOCAB_SIZE} caracteres")
print(f"Caracteres: {''.join(chars[:50])}...")

class CharDataset(Dataset):
    """Dataset de secuencias de caracteres para language modeling."""
    def __init__(self, text, seq_len=128):
        self.data = torch.tensor([char2idx[c] for c in text], dtype=torch.long)
        self.seq_len = seq_len

    def __len__(self):
        return (len(self.data) - 1) // self.seq_len

    def __getitem__(self, idx):
        start = idx * self.seq_len
        x = self.data[start:start + self.seq_len]
        y = self.data[start + 1:start + self.seq_len + 1]  # shifted by 1
        return x, y  # x: input chars, y: next chars (targets)

SEQ_LEN = 128
BATCH_SIZE = 64

# Usar el 90% para train, 10% para val
split = int(len(text) * 0.9)
train_char_ds = CharDataset(text[:split], seq_len=SEQ_LEN)
val_char_ds   = CharDataset(text[split:], seq_len=SEQ_LEN)

train_char_loader = DataLoader(train_char_ds, batch_size=BATCH_SIZE, shuffle=True)
val_char_loader   = DataLoader(val_char_ds, batch_size=BATCH_SIZE, shuffle=False)

print(f"Train sequences: {len(train_char_ds):,}")
print(f"Val sequences:   {len(val_char_ds):,}")

# Ejemplo
x_ex, y_ex = train_char_ds[0]
print(f"\nInput:  {''.join(idx2char[i.item()] for i in x_ex[:40])}…")
print(f"Target: {''.join(idx2char[i.item()] for i in y_ex[:40])}…")

class LSTMCharGen(nn.Module):
    """
    Language model char-level.
    - Entrenamiento: many-to-many (teacher forcing, cada char → next char).
    - Inferencia: one-to-many (seed → generar N chars autoregressivamente).
    """
    def __init__(self, vocab_size, embed_dim=128, hidden_size=256,
                 num_layers=2, dropout=0.2):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
        )
        self.fc = nn.Linear(hidden_size, vocab_size)  # logits sobre el vocabulario
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, hidden=None):
        """
        x: (batch, seq_len) — índices de caracteres
        hidden: (h, c) o None
        Returns: logits (batch, seq_len, vocab_size), (h_n, c_n)
        """
        emb = self.dropout(self.embedding(x))       # (B, T, E)
        out, hidden = self.lstm(emb, hidden)          # (B, T, H)
        logits = self.fc(self.dropout(out))           # (B, T, V)
        return logits, hidden

    def compute_loss(self, batch, criterion):
        x, y = batch
        x, y = x.to(device), y.to(device)
        logits, _ = self.forward(x)
        # CrossEntropy espera (B*T, V) y (B*T,)
        return criterion(logits.view(-1, logits.size(-1)), y.view(-1))

model_otm = LSTMCharGen(
    vocab_size=CHAR_VOCAB_SIZE, embed_dim=128,
    hidden_size=256, num_layers=2
).to(device)
count_params(model_otm)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model_otm.parameters(), lr=2e-3)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

EPOCHS = 20
train_losses, val_losses = [], []

for epoch in range(EPOCHS):
    t_loss = train_epoch(model_otm, train_char_loader, criterion, optimizer)
    v_loss = evaluate(model_otm, val_char_loader, criterion)
    scheduler.step()
    train_losses.append(t_loss)
    val_losses.append(v_loss)

    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1:2d} | Train: {t_loss:.4f} | Val: {v_loss:.4f} | "
              f"Perplexity: {math.exp(v_loss):.1f}")

plot_losses(train_losses, val_losses, "One-to-Many: Char-level LM")

def generate_text(model, seed_text, length=500, temperature=0.8):
    """
    Genera texto de forma autoregresiva (one-to-many).
    1. Procesa seed_text para inicializar hidden state.
    2. Genera 'length' caracteres muestreando de la distribución.
    """
    model.eval()
    with torch.no_grad():
        # Codificar seed
        input_ids = torch.tensor([[char2idx[c] for c in seed_text]],
                                  dtype=torch.long).to(device)

        # Warm-up: procesar seed
        logits, hidden = model(input_ids)

        # Tomar la predicción del último carácter del seed
        last_logits = logits[0, -1, :] / temperature
        probs = torch.softmax(last_logits, dim=0)
        next_char_idx = torch.multinomial(probs, 1).item()

        generated = list(seed_text)
        generated.append(idx2char[next_char_idx])

        # Generar el resto one-to-many
        for _ in range(length - 1):
            x = torch.tensor([[next_char_idx]], dtype=torch.long).to(device)
            logits, hidden = model(x, hidden)  # (1, 1, V)

            last_logits = logits[0, 0, :] / temperature
            probs = torch.softmax(last_logits, dim=0)
            next_char_idx = torch.multinomial(probs, 1).item()
            generated.append(idx2char[next_char_idx])

    return ''.join(generated)

# Generar con diferentes temperaturas
for temp in [0.3, 0.7, 1.0]:
    print(f"\n{'='*60}")
    print(f"🌡️ Temperatura = {temp}")
    print('='*60)
    text = generate_text(model_otm, "ROMEO:", length=200, temperature=temp)
    print(text)

# Dataset sintético de POS tagging en español
# Etiquetas UPOS: DET, NOUN, VERB, ADJ, ADV, PRON, ADP, CONJ, PUNCT
SENTENCES = [
    (["El", "gato", "negro", "duerme", "tranquilamente", "."],
     ["DET", "NOUN", "ADJ", "VERB", "ADV", "PUNCT"]),
    (["La", "casa", "grande", "tiene", "un", "jardín", "bonito", "."],
     ["DET", "NOUN", "ADJ", "VERB", "DET", "NOUN", "ADJ", "PUNCT"]),
    (["Yo", "como", "una", "manzana", "roja", "."],
     ["PRON", "VERB", "DET", "NOUN", "ADJ", "PUNCT"]),
    (["El", "perro", "corre", "rápidamente", "por", "el", "parque", "."],
     ["DET", "NOUN", "VERB", "ADV", "ADP", "DET", "NOUN", "PUNCT"]),
    (["Ella", "lee", "un", "libro", "interesante", "y", "largo", "."],
     ["PRON", "VERB", "DET", "NOUN", "ADJ", "CONJ", "ADJ", "PUNCT"]),
    (["Los", "niños", "juegan", "en", "el", "patio", "."],
     ["DET", "NOUN", "VERB", "ADP", "DET", "NOUN", "PUNCT"]),
    (["Mi", "hermana", "trabaja", "mucho", "."],
     ["DET", "NOUN", "VERB", "ADV", "PUNCT"]),
    (["El", "agua", "fría", "cae", "del", "cielo", "gris", "."],
     ["DET", "NOUN", "ADJ", "VERB", "ADP", "NOUN", "ADJ", "PUNCT"]),
    (["Nosotros", "queremos", "una", "solución", "rápida", "."],
     ["PRON", "VERB", "DET", "NOUN", "ADJ", "PUNCT"]),
    (["Tú", "miras", "la", "televisión", "todas", "las", "noches", "."],
     ["PRON", "VERB", "DET", "NOUN", "DET", "DET", "NOUN", "PUNCT"]),
    # Duplicamos variando para tener más datos
    (["Un", "pájaro", "azul", "vuela", "alto", "."],
     ["DET", "NOUN", "ADJ", "VERB", "ADV", "PUNCT"]),
    (["Las", "flores", "rojas", "crecen", "en", "el", "jardín", "."],
     ["DET", "NOUN", "ADJ", "VERB", "ADP", "DET", "NOUN", "PUNCT"]),
    (["Pedro", "escribe", "cartas", "largas", "y", "bonitas", "."],
     ["NOUN", "VERB", "NOUN", "ADJ", "CONJ", "ADJ", "PUNCT"]),
    (["El", "profesor", "explica", "la", "lección", "claramente", "."],
     ["DET", "NOUN", "VERB", "DET", "NOUN", "ADV", "PUNCT"]),
    (["Ella", "compra", "frutas", "frescas", "en", "el", "mercado", "."],
     ["PRON", "VERB", "NOUN", "ADJ", "ADP", "DET", "NOUN", "PUNCT"]),
]

# Multiplicar datos con variaciones de orden
all_sentences = SENTENCES * 20  # 300 frases para training
np.random.shuffle(all_sentences)

# Vocabulario de palabras y tags
all_words = set()
all_tags = set()
for words, tags in all_sentences:
    all_words.update(w.lower() for w in words)
    all_tags.update(tags)

word2idx_pos = {'<pad>': 0, '<unk>': 1}
for w in sorted(all_words):
    word2idx_pos[w] = len(word2idx_pos)

tag2idx = {'<pad>': 0}
for t in sorted(all_tags):
    tag2idx[t] = len(tag2idx)
idx2tag = {i: t for t, i in tag2idx.items()}

POS_VOCAB_SIZE = len(word2idx_pos)
NUM_TAGS = len(tag2idx)
POS_PAD_IDX = 0
TAG_PAD_IDX = 0

print(f"Vocab: {POS_VOCAB_SIZE} palabras | Tags: {NUM_TAGS}")
print(f"Tags: {list(tag2idx.keys())}")

class POSDataset(Dataset):
    def __init__(self, data, max_len=12):
        self.max_len = max_len
        self.x, self.y = [], []
        for words, tags in data:
            x = [word2idx_pos.get(w.lower(), 1) for w in words][:max_len]
            y = [tag2idx[t] for t in tags][:max_len]
            # Padding
            pad_len = max_len - len(x)
            x += [POS_PAD_IDX] * pad_len
            y += [TAG_PAD_IDX] * pad_len
            self.x.append(x)
            self.y.append(y)
        self.x = torch.tensor(self.x, dtype=torch.long)
        self.y = torch.tensor(self.y, dtype=torch.long)

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

MAX_POS_LEN = 12

# Split: 80% train, 20% val
split_idx = int(len(all_sentences) * 0.8)
train_pos_ds = POSDataset(all_sentences[:split_idx], max_len=MAX_POS_LEN)
val_pos_ds   = POSDataset(all_sentences[split_idx:], max_len=MAX_POS_LEN)

train_pos_loader = DataLoader(train_pos_ds, batch_size=32, shuffle=True)
val_pos_loader   = DataLoader(val_pos_ds, batch_size=32, shuffle=False)

print(f"Train: {len(train_pos_ds)} | Val: {len(val_pos_ds)}")
print(f"Input shape:  {train_pos_ds[0][0].shape}")   # (12,)
print(f"Target shape: {train_pos_ds[0][1].shape}")    # (12,)

class LSTMPOSTagger(nn.Module):
    """
    Many-to-Many síncrono: BiLSTM → Linear en cada timestep.
    Cada palabra recibe una etiqueta POS.
    """
    def __init__(self, vocab_size, num_tags, embed_dim=64, hidden_size=64,
                 num_layers=1, dropout=0.2, pad_idx=0):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,  # Bidireccional: crucial para tagging
            dropout=dropout if num_layers > 1 else 0,
        )
        self.dropout = nn.Dropout(dropout)
        # FC: para cada timestep, predecir el tag
        self.fc = nn.Linear(hidden_size * 2, num_tags)  # *2 por bidireccional
        self.tag_pad_idx = 0

    def forward(self, x):
        """
        x: (batch, seq_len) — índices de palabras
        Returns: logits (batch, seq_len, num_tags)
        """
        emb = self.dropout(self.embedding(x))      # (B, T, E)
        output, _ = self.lstm(emb)                   # (B, T, H*2)
        logits = self.fc(self.dropout(output))       # (B, T, num_tags)
        return logits

    def compute_loss(self, batch, criterion):
        x, y = batch
        x, y = x.to(device), y.to(device)
        logits = self.forward(x)         # (B, T, C)
        # Flatten: (B*T, C) y (B*T,)
        logits_flat = logits.view(-1, logits.size(-1))
        y_flat = y.view(-1)
        return criterion(logits_flat, y_flat)

model_mtms = LSTMPOSTagger(
    vocab_size=POS_VOCAB_SIZE, num_tags=NUM_TAGS,
    embed_dim=64, hidden_size=64
).to(device)
count_params(model_mtms)

# Ignorar padding en la loss
criterion = nn.CrossEntropyLoss(ignore_index=TAG_PAD_IDX)
optimizer = optim.Adam(model_mtms.parameters(), lr=1e-2)

EPOCHS = 50
train_losses, val_losses = [], []

for epoch in range(EPOCHS):
    t_loss = train_epoch(model_mtms, train_pos_loader, criterion, optimizer)
    v_loss = evaluate(model_mtms, val_pos_loader, criterion)
    train_losses.append(t_loss)
    val_losses.append(v_loss)

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1:2d} | Train: {t_loss:.4f} | Val: {v_loss:.4f}")

plot_losses(train_losses, val_losses, "Many-to-Many Sync: POS Tagging")

# ── Token-level accuracy ──
model_mtms.eval()
correct, total = 0, 0
with torch.no_grad():
    for x, y in val_pos_loader:
        x, y = x.to(device), y.to(device)
        logits = model_mtms(x)       # (B, T, C)
        preds = logits.argmax(dim=-1) # (B, T)
        # Solo contar tokens no-padding
        mask = y != TAG_PAD_IDX
        correct += ((preds == y) & mask).sum().item()
        total += mask.sum().item()

print(f"\n🎯 Token Accuracy (val): {correct/total:.2%}")

# ── Ejemplo de predicción ──
test_sentence = ["El", "perro", "negro", "corre", "rápidamente", "."]
x_test = torch.tensor([[word2idx_pos.get(w.lower(), 1) for w in test_sentence]
                        + [0] * (MAX_POS_LEN - len(test_sentence))],
                       dtype=torch.long).to(device)

with torch.no_grad():
    logits = model_mtms(x_test)
    pred_tags = logits[0].argmax(dim=-1)

print("\nEjemplo de predicción:")
for word, tag_idx in zip(test_sentence, pred_tags[:len(test_sentence)]):
    print(f"  {word:>15s}  →  {idx2tag[tag_idx.item()]}")

import random
from datetime import datetime, timedelta

MONTHS = [
    "January", "February", "March", "April", "May", "June",
    "July", "August", "September", "October", "November", "December"
]

def random_date():
    """Genera una fecha aleatoria entre 1950-2030."""
    start = datetime(1950, 1, 1)
    end = datetime(2030, 12, 31)
    delta = end - start
    random_days = random.randint(0, delta.days)
    return start + timedelta(days=random_days)

def date_to_input(dt):
    """Fecha → formato numérico: '15/03/2024'"""
    return dt.strftime("%d/%m/%Y")

def date_to_output(dt):
    """Fecha → formato texto: 'March 15, 2024'"""
    return f"{MONTHS[dt.month-1]} {dt.day}, {dt.year}"

# Generar datos
random.seed(SEED)
N_DATES = 5000
dates = [random_date() for _ in range(N_DATES)]
pairs = [(date_to_input(d), date_to_output(d)) for d in dates]

# Verificar
for i in range(5):
    print(f"  {pairs[i][0]:>12s}  →  {pairs[i][1]}")

# Vocabulario a nivel de carácter (entrada y salida)
src_chars = sorted(set(''.join(p[0] for p in pairs)))
tgt_chars = sorted(set(''.join(p[1] for p in pairs)))
all_chars_s2s = sorted(set(src_chars + tgt_chars))

SOS_TOKEN, EOS_TOKEN, PAD_TOKEN_S2S = '<S>', '<E>', '<P>'
s2s_vocab = [PAD_TOKEN_S2S, SOS_TOKEN, EOS_TOKEN] + all_chars_s2s

s2s_char2idx = {c: i for i, c in enumerate(s2s_vocab)}
s2s_idx2char = {i: c for c, i in s2s_char2idx.items()}
S2S_VOC_SIZE = len(s2s_vocab)
SOS_IDX = s2s_char2idx[SOS_TOKEN]
EOS_IDX = s2s_char2idx[EOS_TOKEN]
S2S_PAD_IDX = s2s_char2idx[PAD_TOKEN_S2S]

print(f"\nVocab size: {S2S_VOC_SIZE}")
print(f"SOS={SOS_IDX}, EOS={EOS_IDX}, PAD={S2S_PAD_IDX}")

class DateSeq2SeqDataset(Dataset):
    def __init__(self, pairs, max_src_len=12, max_tgt_len=22):
        self.data = []
        for src, tgt in pairs:
            # Codificar source
            src_ids = [s2s_char2idx[c] for c in src]
            src_ids += [S2S_PAD_IDX] * (max_src_len - len(src_ids))

            # Codificar target con SOS y EOS
            tgt_ids = [SOS_IDX] + [s2s_char2idx[c] for c in tgt] + [EOS_IDX]
            tgt_ids += [S2S_PAD_IDX] * (max_tgt_len - len(tgt_ids))

            self.data.append((
                torch.tensor(src_ids[:max_src_len], dtype=torch.long),
                torch.tensor(tgt_ids[:max_tgt_len], dtype=torch.long),
            ))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

MAX_SRC = 12   # "DD/MM/YYYY" = 10 chars + padding
MAX_TGT = 22   # "September 30, 2024" + SOS + EOS + padding

# Split
train_pairs = pairs[:4000]
val_pairs = pairs[4000:4500]
test_pairs = pairs[4500:]

train_s2s_ds = DateSeq2SeqDataset(train_pairs, MAX_SRC, MAX_TGT)
val_s2s_ds   = DateSeq2SeqDataset(val_pairs, MAX_SRC, MAX_TGT)
test_s2s_ds  = DateSeq2SeqDataset(test_pairs, MAX_SRC, MAX_TGT)

train_s2s_loader = DataLoader(train_s2s_ds, batch_size=64, shuffle=True)
val_s2s_loader   = DataLoader(val_s2s_ds, batch_size=64, shuffle=False)
test_s2s_loader  = DataLoader(test_s2s_ds, batch_size=64, shuffle=False)

print(f"Train: {len(train_s2s_ds)} | Val: {len(val_s2s_ds)} | Test: {len(test_s2s_ds)}")
src_ex, tgt_ex = train_s2s_ds[0]
print(f"Source shape: {src_ex.shape} | Target shape: {tgt_ex.shape}")

class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_layers=1, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=S2S_PAD_IDX)
        self.lstm = nn.LSTM(embed_dim, hidden_size, num_layers,
                            batch_first=True, dropout=dropout if num_layers > 1 else 0)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        """src: (B, T_src) → hidden: (h_n, c_n)"""
        emb = self.dropout(self.embedding(src))
        _, hidden = self.lstm(emb)
        return hidden  # (h_n, c_n) se pasan al decoder


class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_size, num_layers=1, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=S2S_PAD_IDX)
        self.lstm = nn.LSTM(embed_dim, hidden_size, num_layers,
                            batch_first=True, dropout=dropout if num_layers > 1 else 0)
        self.fc = nn.Linear(hidden_size, vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, tgt, hidden):
        """
        tgt: (B, T_tgt) — secuencia target (con SOS al inicio)
        hidden: (h_n, c_n) del encoder
        Returns: logits (B, T_tgt, vocab_size)
        """
        emb = self.dropout(self.embedding(tgt))
        output, hidden = self.lstm(emb, hidden)
        logits = self.fc(output)
        return logits, hidden

    def decode_step(self, x, hidden):
        """Un paso de decodificación: x: (B, 1) → logits: (B, 1, V)"""
        emb = self.dropout(self.embedding(x))
        output, hidden = self.lstm(emb, hidden)
        logits = self.fc(output)
        return logits, hidden


class Seq2Seq(nn.Module):
    """
    Encoder-Decoder Seq2Seq con LSTM.
    - Entrenamiento: teacher forcing (el decoder recibe el target real).
    - Inferencia: autoregresivo (el decoder recibe su propia predicción).
    """
    def __init__(self, vocab_size, embed_dim=64, hidden_size=128, num_layers=2, dropout=0.1):
        super().__init__()
        self.encoder = Encoder(vocab_size, embed_dim, hidden_size, num_layers, dropout)
        self.decoder = Decoder(vocab_size, embed_dim, hidden_size, num_layers, dropout)
        self.vocab_size = vocab_size

    def forward(self, src, tgt):
        """
        Teacher forcing: src → encoder → hidden → decoder(tgt) → logits
        """
        hidden = self.encoder(src)
        # tgt[:, :-1]: todo excepto el último token (el decoder no necesita predecir después de EOS)
        logits, _ = self.decoder(tgt[:, :-1], hidden)
        return logits  # (B, T_tgt-1, V)

    def compute_loss(self, batch, criterion):
        src, tgt = batch
        src, tgt = src.to(device), tgt.to(device)
        logits = self.forward(src, tgt)
        # Target: tgt[:, 1:] (sin el SOS)
        target = tgt[:, 1:logits.size(1) + 1]
        return criterion(logits.reshape(-1, self.vocab_size), target.reshape(-1))

    def translate(self, src, max_len=22):
        """Inferencia autoregresiva."""
        self.eval()
        with torch.no_grad():
            hidden = self.encoder(src)
            batch_size = src.size(0)
            # Empezar con SOS
            input_tok = torch.full((batch_size, 1), SOS_IDX, dtype=torch.long, device=src.device)
            outputs = []
            for _ in range(max_len):
                logits, hidden = self.decoder.decode_step(input_tok, hidden)
                next_tok = logits.argmax(dim=-1)  # (B, 1)
                outputs.append(next_tok)
                input_tok = next_tok
            return torch.cat(outputs, dim=1)  # (B, max_len)

model_s2s = Seq2Seq(
    vocab_size=S2S_VOC_SIZE, embed_dim=64,
    hidden_size=128, num_layers=2
).to(device)
count_params(model_s2s)

criterion = nn.CrossEntropyLoss(ignore_index=S2S_PAD_IDX)
optimizer = optim.Adam(model_s2s.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)

EPOCHS = 30
train_losses, val_losses = [], []

for epoch in range(EPOCHS):
    t_loss = train_epoch(model_s2s, train_s2s_loader, criterion, optimizer)
    v_loss = evaluate(model_s2s, val_s2s_loader, criterion)
    scheduler.step(v_loss)
    train_losses.append(t_loss)
    val_losses.append(v_loss)

    if (epoch + 1) % 5 == 0:
        print(f"Epoch {epoch+1:2d} | Train: {t_loss:.4f} | Val: {v_loss:.4f}")

plot_losses(train_losses, val_losses, "Many-to-Many Async: Seq2Seq Date Conversion")

def decode_output(indices):
    """Convierte índices a string, cortando en EOS."""
    chars = []
    for idx in indices:
        if idx == EOS_IDX:
            break
        if idx not in (S2S_PAD_IDX, SOS_IDX):
            chars.append(s2s_idx2char.get(idx, '?'))
    return ''.join(chars)

# Evaluar en test
model_s2s.eval()
correct, total = 0, 0
examples = []

for src, tgt in test_s2s_loader:
    src = src.to(device)
    pred_indices = model_s2s.translate(src, max_len=MAX_TGT)

    for i in range(src.size(0)):
        # Decodificar predicción
        pred_str = decode_output(pred_indices[i].cpu().tolist())
        # Decodificar target
        tgt_str = decode_output(tgt[i].tolist())

        if pred_str == tgt_str:
            correct += 1
        total += 1

        if len(examples) < 10:
            src_str = decode_output(src[i].cpu().tolist())
            examples.append((src_str, tgt_str, pred_str, pred_str == tgt_str))

print(f"🎯 Exact Match Accuracy (test): {correct/total:.2%}\n")
print("Ejemplos de traducción:")
print(f"  {'INPUT':>14s}  {'EXPECTED':>22s}  {'PREDICTED':>22s}  {'OK':>4s}")
print(f"  {'─'*14}  {'─'*22}  {'─'*22}  {'─'*4}")
for src_s, tgt_s, pred_s, ok in examples:
    mark = '✓' if ok else '✗'
    print(f"  {src_s:>14s}  {tgt_s:>22s}  {pred_s:>22s}  {mark:>4s}")