import torch

# The shape challenge in Seq2Seq:
# Input and output have DIFFERENT lengths → you can't use a simple many-to-many RNN

# Example: English → French translation
src_sentences = [
    "the cat sat on the mat",        # 6 words
    "I love machine learning",       # 4 words
    "natural language processing is fascinating",  # 5 words
]

tgt_sentences = [
    "le chat était assis sur le tapis",   # 7 words (different length!)
    "j'adore l'apprentissage automatique", # 2 words!
    "le traitement du langage naturel est fascinant",  # 8 words
]

# A standard RNN's hidden state at the end of the source sequence is a fixed-size vector.
# The decoder must generate a VARIABLE-length output from this fixed-size starting point.
# This is the fundamental challenge that the Encoder-Decoder architecture addresses.

print("Seq2Seq length mismatch examples:")
for src, tgt in zip(src_sentences, tgt_sentences):
    print(f"  Source ({len(src.split())} words): '{src}'")
    print(f"  Target ({len(tgt.split())} words): '{tgt}'")
    print()

import torch
import torch.nn as nn

class Encoder(nn.Module):
    """
    Bidirectional LSTM encoder.
    Reads the entire source sequence and produces a context vector.
    """
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers=2, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=True   # read forwards AND backwards
        )
        # Project from 2*hidden_dim (bidirectional) to hidden_dim for decoder
        self.fc_h = nn.Linear(hidden_dim * 2, hidden_dim)
        self.fc_c = nn.Linear(hidden_dim * 2, hidden_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, src_lengths):
        """
        src:         (batch, src_len)    — source token IDs
        src_lengths: (batch,)            — actual lengths (for padding mask)

        Returns:
            encoder_outputs: (batch, src_len, 2*hidden_dim) — all hidden states
            hidden:          (batch, hidden_dim)            — context vector (h)
            cell:            (batch, hidden_dim)            — context vector (c)
        """
        embedded = self.dropout(self.embedding(src))  # (B, T, E)

        # Pack to skip padding positions
        packed = nn.utils.rnn.pack_padded_sequence(
            embedded, src_lengths.cpu(), batch_first=True, enforce_sorted=False
        )
        packed_outputs, (h_n, c_n) = self.lstm(packed)

        # Unpack outputs: (B, T, 2*hidden_dim)
        outputs, _ = nn.utils.rnn.pad_packed_sequence(packed_outputs, batch_first=True)

        # h_n shape: (num_layers * 2, B, hidden_dim)
        # Take last layer's forward (h_n[-2]) and backward (h_n[-1]) directions
        h_forward  = h_n[-2, :, :]   # (B, hidden_dim) — last layer forward
        h_backward = h_n[-1, :, :]   # (B, hidden_dim) — last layer backward
        c_forward  = c_n[-2, :, :]
        c_backward = c_n[-1, :, :]

        # Concatenate and project to single hidden_dim for the decoder
        hidden = torch.tanh(self.fc_h(torch.cat([h_forward, h_backward], dim=1)))
        cell   = torch.tanh(self.fc_c(torch.cat([c_forward, c_backward], dim=1)))

        return outputs, hidden, cell


# Quick test
encoder = Encoder(vocab_size=5000, embed_dim=256, hidden_dim=512, num_layers=2)
src = torch.randint(1, 5000, (4, 12))    # batch=4, max_src_len=12
src_lengths = torch.tensor([12, 10, 8, 6])

enc_out, h, c = encoder(src, src_lengths)
print(f"Encoder outputs shape:  {enc_out.shape}")  # (4, 12, 1024) — 2*512
print(f"Context vector h shape: {h.shape}")        # (4, 512)
print(f"Context vector c shape: {c.shape}")        # (4, 512)

import torch
import torch.nn as nn
import torch.nn.functional as F

class Decoder(nn.Module):
    """
    LSTM decoder that generates one token per step.
    """
    def __init__(self, vocab_size, embed_dim, hidden_dim, dropout=0.3):
        super().__init__()
        self.embedding  = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        # Note: unidirectional! Can't look ahead during generation
        self.lstm = nn.LSTMCell(
            input_size=embed_dim,
            hidden_size=hidden_dim,
        )
        self.fc_out = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward_step(self, token_ids, h, c):
        """
        Single decoding step — predict next token.

        token_ids: (batch,)       — input token IDs for this step
        h:         (batch, H)     — hidden state
        c:         (batch, H)     — cell state

        Returns:
            logits: (batch, vocab_size) — unnormalized scores for each token
            h_new:  (batch, H)
            c_new:  (batch, H)
        """
        embedded = self.dropout(self.embedding(token_ids))  # (B, E)
        h_new, c_new = self.lstm(embedded, (h, c))          # LSTMCell: one step
        logits = self.fc_out(h_new)                         # (B, vocab_size)
        return logits, h_new, c_new

    def greedy_decode(self, h, c, sos_id, eos_id, max_len=50):
        """
        Greedy decoding: always pick the most likely next token.
        Stops when EOS is generated or max_len is reached.
        """
        batch_size = h.size(0)
        # Start with SOS token
        current_token = torch.full((batch_size,), sos_id, dtype=torch.long)
        generated = []
        finished  = torch.zeros(batch_size, dtype=torch.bool)

        for _ in range(max_len):
            logits, h, c = self.forward_step(current_token, h, c)
            current_token = logits.argmax(dim=-1)  # greedy: pick best
            generated.append(current_token)
            finished |= (current_token == eos_id)
            if finished.all():
                break

        return torch.stack(generated, dim=1)  # (batch, gen_len)


# Quick test
decoder = Decoder(vocab_size=5000, embed_dim=256, hidden_dim=512)
h0 = torch.randn(4, 512)   # context vector from encoder (4 examples)
c0 = torch.randn(4, 512)

SOS_ID, EOS_ID = 1, 2
generated = decoder.greedy_decode(h0, c0, sos_id=SOS_ID, eos_id=EOS_ID, max_len=20)
print(f"Generated sequence shape: {generated.shape}")  # (4, ≤20)

import torch
import torch.nn as nn
import random

def train_step_with_teacher_forcing(encoder, decoder, src, src_lengths,
                                     tgt, teacher_forcing_ratio=0.5):
    """
    Training step with scheduled teacher forcing.

    teacher_forcing_ratio: probability of using ground truth token
                          (0.0 = always use prediction, 1.0 = always use truth)
    """
    # Encode source sequence
    enc_outputs, h, c = encoder(src, src_lengths)

    batch_size = src.size(0)
    tgt_len    = tgt.size(1)
    vocab_size = decoder.fc_out.out_features

    # Tensor to hold all decoder predictions
    predictions = torch.zeros(batch_size, tgt_len, vocab_size)

    # First decoder input: SOS token
    dec_input = tgt[:, 0]  # (batch,) — SOS token at position 0

    for t in range(1, tgt_len):
        logits, h, c = decoder.forward_step(dec_input, h, c)
        predictions[:, t, :] = logits

        # Teacher forcing decision
        if random.random() < teacher_forcing_ratio:
            dec_input = tgt[:, t]          # ← use GROUND TRUTH token
        else:
            dec_input = logits.argmax(1)   # ← use MODEL PREDICTION

    return predictions


# Why teacher forcing is problematic:
print("Teacher Forcing — Training vs Inference mismatch:")
print()
print("  During TRAINING with teacher forcing:")
print("    Step 1: input=SOS   → model predicts 'le'  (correct!)")
print("    Step 2: input='le'  → model predicts 'chat' (correct!)  [teacher forced]")
print("    Step 3: input='chat'→ model predicts 'était' (correct!) [teacher forced]")
print()
print("  During INFERENCE (no teacher forcing available):")
print("    Step 1: input=SOS   → model predicts 'le'   (correct!)")
print("    Step 2: input='le'  → model predicts 'les'  (WRONG! should be 'chat')")
print("    Step 3: input='les' → model has NEVER seen 'les' as context in training!")
print("    → Error cascades catastrophically!")
print()
print("  Solution: scheduled teacher forcing — start at ratio=1.0,")
print("  gradually decrease to 0.5 as training progresses.")

import torch

def beam_search(decoder, h_init, c_init, sos_id, eos_id,
                max_len=50, beam_size=4, length_penalty=0.7):
    """
    Beam search decoding — maintain top-K candidate sequences.

    At each step:
    1. Expand each of the K beams by generating all possible next tokens
    2. Score each expanded sequence: log(P(sequence)) = sum of log-probs
    3. Keep only the top-K scoring sequences
    4. Stop when all K beams have generated EOS or max_len is reached
    """
    device = h_init.device
    # Start with single SOS token in each beam
    beams = [(0.0, [sos_id], h_init, c_init)]  # (score, tokens, h, c)
    completed = []

    for step in range(max_len):
        new_beams = []

        for score, tokens, h, c in beams:
            # Stop expanding this beam if it ended
            if tokens[-1] == eos_id:
                completed.append((score, tokens))
                continue

            # Decode one step
            current = torch.tensor([tokens[-1]], device=device)
            logits, h_new, c_new = decoder.forward_step(current, h, c)
            log_probs = torch.log_softmax(logits[0], dim=-1)

            # Get top beam_size next tokens for this beam
            top_log_probs, top_ids = log_probs.topk(beam_size)

            for log_p, token_id in zip(top_log_probs, top_ids):
                new_score  = score + log_p.item()
                new_tokens = tokens + [token_id.item()]
                new_beams.append((new_score, new_tokens, h_new, c_new))

        # Keep only top beam_size beams (sorted by score)
        # Length penalty: penalize shorter sequences so they don't dominate
        def score_fn(beam):
            score, tokens, _, _ = beam
            length = len(tokens)
            lp = ((5 + length) / 6) ** length_penalty  # standard length penalty
            return score / lp

        beams = sorted(new_beams, key=score_fn, reverse=True)[:beam_size]

        if not beams:
            break

    # Return best completed beam (or best active beam if none completed)
    if completed:
        best = max(completed, key=lambda x: x[0] / len(x[1]))
    else:
        best = max(beams, key=lambda x: score_fn(x))[:2]
    return best[1]


# Illustrate the benefit of beam search:
print("Greedy vs Beam Search:")
print()
print("Example: Translating 'I have a cold'")
print()
print("Step-by-step probabilities:")
print("  Greedy: always picks max at each step")
print("    P('I')        = 0.9  → picks 'I'")
print("    P('am'|I)     = 0.7  → picks 'am'  ← commits early!")
print("    P('cold'|I am)= 0.6  → 'I am cold' (wrong meaning)")
print()
print("  Beam Search (K=3): keeps top 3 sequences")
print("    Step 1: ['I'(0.9), 'I'(0.9), 'He'(0.05)]")
print("    Step 2: ['I am'(0.63), 'I have'(0.27), 'I feel'(0.18)]")
print("    Step 3: ['I am cold'(0.38), 'I have a cold'(0.24), 'I feel cold'(0.15)]")
print("    After length norm: 'I have a cold' scores highest!")
print()
print("Typical beam_size values in production MT:")
print("  beam_size=4:  good quality/speed tradeoff (default in many systems)")
print("  beam_size=10: better quality, 2-3x slower")
print("  beam_size=1:  equivalent to greedy decoding")

# Using Hugging Face transformers for beam search (real-world usage)
from transformers import T5ForConditionalGeneration, T5Tokenizer

model     = T5ForConditionalGeneration.from_pretrained('t5-small')
tokenizer = T5Tokenizer.from_pretrained('t5-small')

text = "translate English to French: The house is wonderful."
inputs = tokenizer(text, return_tensors='pt')

# Greedy decoding (beam_size=1)
greedy_output = model.generate(
    **inputs,
    max_length=40,
    num_beams=1,           # ← greedy
)

# Beam search decoding (beam_size=4)
beam_output = model.generate(
    **inputs,
    max_length=40,
    num_beams=4,           # ← beam search
    early_stopping=True,
    length_penalty=0.6,
    no_repeat_ngram_size=2 # prevent repetition
)

print("Input:  ", text)
print("Greedy: ", tokenizer.decode(greedy_output[0], skip_special_tokens=True))
print("Beam-4: ", tokenizer.decode(beam_output[0], skip_special_tokens=True))

from nltk.translate.bleu_score import sentence_bleu, corpus_bleu, SmoothingFunction
import nltk
nltk.download('punkt', quiet=True)

# BLEU score on sentence level
reference  = [["the", "cat", "is", "on", "the", "mat"]]   # list of references
hypothesis1 = ["the", "cat", "is", "on", "the", "mat"]    # perfect match
hypothesis2 = ["the", "cat", "sat", "on", "the", "mat"]   # one word different
hypothesis3 = ["a", "cat", "stood", "on", "a", "carpet"]  # quite different
hypothesis4 = ["the", "mat"]                               # too short (brevity penalty)

smooth = SmoothingFunction().method1  # avoids zero scores for no n-gram matches

for i, hyp in enumerate([hypothesis1, hypothesis2, hypothesis3, hypothesis4], 1):
    bleu1 = sentence_bleu(reference, hyp, weights=(1,0,0,0))          # 1-gram
    bleu4 = sentence_bleu(reference, hyp, weights=(0.25,)*4, smoothing_function=smooth)
    print(f"Hyp {i}: BLEU-1={bleu1:.3f}, BLEU-4={bleu4:.3f}  {hyp}")

# Hyp 1: BLEU-1=1.000, BLEU-4=1.000  (perfect)
# Hyp 2: BLEU-1=0.833, BLEU-4=0.672  (close)
# Hyp 3: BLEU-1=0.333, BLEU-4=0.113  (quite different)
# Hyp 4: BLEU-1=1.000, BLEU-4=0.052  (short: 1-gram precision is high but brevity penalty kills BLEU-4)

# BLEU formula breakdown
import math

def manual_bleu4(reference_tokens, hypothesis_tokens):
    """
    Manual BLEU-4 calculation to show each component.
    """
    ref  = reference_tokens
    hyp  = hypothesis_tokens

    precision_scores = []

    for n in range(1, 5):
        # Count n-grams in hypothesis
        hyp_ngrams = {}
        for i in range(len(hyp) - n + 1):
            gram = tuple(hyp[i:i+n])
            hyp_ngrams[gram] = hyp_ngrams.get(gram, 0) + 1

        # Count n-grams in reference
        ref_ngrams = {}
        for i in range(len(ref) - n + 1):
            gram = tuple(ref[i:i+n])
            ref_ngrams[gram] = ref_ngrams.get(gram, 0) + 1

        # Clipped count: min(hypothesis count, reference count) for each n-gram
        clipped = sum(min(count, ref_ngrams.get(gram, 0))
                      for gram, count in hyp_ngrams.items())

        total_hyp_ngrams = max(len(hyp) - n + 1, 1)
        prec = clipped / total_hyp_ngrams
        precision_scores.append(prec)
        print(f"  {n}-gram precision: {clipped}/{total_hyp_ngrams} = {prec:.4f}")

    # Brevity penalty: punish translations that are too short
    bp = min(1.0, math.exp(1 - len(ref) / max(len(hyp), 1)))
    print(f"  Brevity penalty: {bp:.4f}")

    # BLEU-4: geometric mean of 1-4 gram precisions × brevity penalty
    # Avoid log(0) by smoothing
    log_avg = sum(math.log(p + 1e-10) for p in precision_scores) / 4
    bleu = bp * math.exp(log_avg)
    print(f"  BLEU-4: {bleu:.4f}")
    return bleu

print("Reference: 'the cat sat on the mat'")
print("Hypothesis: 'the cat is on the mat'")
print()
manual_bleu4(
    ["the", "cat", "sat", "on", "the", "mat"],
    ["the", "cat", "is",  "on", "the", "mat"]
)

import torch
import torch.nn as nn
import numpy as np

# Simulate the information bottleneck:
# As sequence length grows, the final hidden state must represent more information
# in the same fixed-size vector

class SimpleEncoder(nn.Module):
    def __init__(self, input_size=10, hidden_size=32):
        super().__init__()
        self.rnn = nn.LSTM(input_size, hidden_size, batch_first=True)

    def forward(self, x):
        _, (h_n, _) = self.rnn(x)
        return h_n.squeeze(0)  # final hidden state = context vector

encoder = SimpleEncoder()

# Test: can the encoder preserve information about the FIRST word
# across sequences of different lengths?
def test_first_word_memory(seq_lengths, hidden_size=32, trials=100):
    """
    For each sequence length: how different is the context vector
    when we change the FIRST word vs changing a LATE word?
    A good encoder should produce different vectors for both changes.
    A bottlenecked encoder loses information about early words.
    """
    results = {}
    for seq_len in seq_lengths:
        first_word_diffs = []
        for _ in range(trials):
            # Create two sequences identical except for word at position 0
            seq_a = torch.randn(1, seq_len, 10)
            seq_b = seq_a.clone()
            seq_b[0, 0, :] += 2.0  # change only first word

            ctx_a = encoder(seq_a).detach().numpy()
            ctx_b = encoder(seq_b).detach().numpy()

            diff = np.linalg.norm(ctx_a - ctx_b)
            first_word_diffs.append(diff)

        results[seq_len] = np.mean(first_word_diffs)

    return results

seq_lengths = [5, 10, 20, 50, 100]
print("Information bottleneck demonstration:")
print("Sequence length | Context diff (first word change) | Interpretation")
print("-" * 70)

diffs = test_first_word_memory(seq_lengths)
for length, diff in diffs.items():
    quality = "✓ Good" if diff > 0.3 else "⚠ Weak" if diff > 0.05 else "✗ Lost"
    print(f"  Length {length:5d}   |   {diff:.4f}                            | {quality}")

import matplotlib.pyplot as plt

# Empirical results from Sutskever et al. 2014 (the original Seq2Seq paper)
# BLEU scores on English-French translation by sentence length bucket
sentence_lengths = [10, 15, 20, 25, 30, 35, 40, 45, 50]
bleu_without_attention = [33.1, 31.5, 29.8, 27.2, 23.1, 19.4, 15.2, 11.8, 9.3]
bleu_with_attention    = [35.2, 34.1, 33.8, 32.9, 31.7, 30.5, 29.2, 28.1, 26.8]

plt.figure(figsize=(9, 5))
plt.plot(sentence_lengths, bleu_without_attention, 'r-o', label='Seq2Seq (no attention)')
plt.plot(sentence_lengths, bleu_with_attention,    'g-o', label='Seq2Seq + Attention (Bahdanau 2015)')
plt.axhline(30, color='gray', linestyle='--', alpha=0.5, label='BLEU=30 (good MT quality)')
plt.xlabel('Source sentence length (words)')
plt.ylabel('BLEU score')
plt.title('BLEU Score Degrades with Sequence Length — Fixed by Attention')
plt.legend()
plt.tight_layout()
plt.show()

print("Key observation:")
print("  Without attention: BLEU drops from 33 to 9 as length grows 10 → 50 words")
print("  With attention:    BLEU stays ~29-35 regardless of sentence length!")
print()
print("This plot is what motivated Bahdanau et al. to invent attention in 2015.")

import torch
import torch.nn as nn
import random

class Seq2Seq(nn.Module):
    """
    Complete Encoder-Decoder Seq2Seq model.
    """
    def __init__(self, encoder, decoder, pad_idx, device):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.pad_idx = pad_idx
        self.device  = device

    def forward(self, src, src_lengths, tgt, teacher_forcing_ratio=0.5):
        """
        src:          (batch, src_len)   — source token IDs
        src_lengths:  (batch,)           — source lengths
        tgt:          (batch, tgt_len)   — target token IDs (includes SOS, EOS)
        Returns: predictions (batch, tgt_len, vocab_size)
        """
        batch_size  = src.size(0)
        tgt_len     = tgt.size(1)
        vocab_size  = self.decoder.fc_out.out_features

        # Storage for all decoder outputs
        outputs = torch.zeros(batch_size, tgt_len, vocab_size).to(self.device)

        # Encode the source sequence → context vector
        enc_out, h, c = self.encoder(src, src_lengths)

        # Start decoder with SOS token (tgt[:,0] = SOS for all in batch)
        dec_input = tgt[:, 0]

        for t in range(1, tgt_len):
            # Single decoder step
            logits, h, c = self.decoder.forward_step(dec_input, h, c)
            outputs[:, t, :] = logits

            # Teacher forcing: use true label or model prediction?
            teacher_force = random.random() < teacher_forcing_ratio
            dec_input = tgt[:, t] if teacher_force else logits.argmax(1)

        return outputs


# Training recipe
def create_model(src_vocab_size, tgt_vocab_size, embed_dim=256,
                 hidden_dim=512, num_layers=2, dropout=0.3, device='cpu'):
    enc = Encoder(src_vocab_size, embed_dim, hidden_dim, num_layers, dropout)
    dec = Decoder(tgt_vocab_size, embed_dim, hidden_dim, dropout)
    model = Seq2Seq(enc, dec, pad_idx=0, device=device)
    return model


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

model = create_model(src_vocab_size=8000, tgt_vocab_size=10000)
print(f"Total trainable parameters: {count_parameters(model):,}")

# Training loop sketch
import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=5e-4, weight_decay=1e-5)
criterion = nn.CrossEntropyLoss(ignore_index=0)  # ignore PAD tokens
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)

# Standard training step:
def train_batch(model, batch, optimizer, criterion, clip=1.0):
    model.train()
    optimizer.zero_grad()

    src, src_len, tgt = batch
    # tgt: (batch, tgt_len) — includes SOS and EOS

    output = model(src, src_len, tgt, teacher_forcing_ratio=0.5)
    # output: (batch, tgt_len, vocab_size)

    # Flatten for cross-entropy: skip position 0 (SOS)
    output_flat = output[:, 1:].reshape(-1, output.size(-1))
    target_flat = tgt[:, 1:].reshape(-1)

    loss = criterion(output_flat, target_flat)
    loss.backward()

    # IMPORTANT: gradient clipping for RNN stability
    torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
    optimizer.step()

    return loss.item()

print("\nTraining setup:")
print(f"  Optimizer: Adam, lr=5e-4")
print(f"  Loss: CrossEntropy (PAD tokens masked)")
print(f"  Gradient clipping: max_norm=1.0")
print(f"  Teacher forcing: 0.5 (scheduled down during training)")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import random

# ── Dataset ──
SRC_VOCAB = {'0':2, '1':3, '2':4, '3':5, '4':6,
             '5':7, '6':8, '7':9, '8':10, '9':11,
             '<PAD>':0, '<SOS>':1, '<EOS>':12}

TGT_VOCAB = {'zero':2, 'one':3, 'two':4, 'three':5, 'four':6,
             'five':7, 'six':8, 'seven':9, 'eight':10, 'nine':11,
             '<PAD>':0, '<SOS>':1, '<EOS>':12}

TGT_IDX_TO_WORD = {v: k for k, v in TGT_VOCAB.items()}
WORDS = ['zero','one','two','three','four','five','six','seven','eight','nine']

def make_pair(length=3):
    """Generate a random number → words pair."""
    digits = [random.randint(0, 9) for _ in range(length)]
    src_ids = [SRC_VOCAB[str(d)] for d in digits]
    tgt_ids = [TGT_VOCAB['<SOS>']] + [TGT_VOCAB[WORDS[d]] for d in digits] + [TGT_VOCAB['<EOS>']]
    return digits, src_ids, tgt_ids

# Demonstrate the task
print("Sample training pairs (digit seq → word seq):")
for _ in range(5):
    digits, src, tgt = make_pair(3)
    src_str = ''.join(str(d) for d in digits)
    tgt_str = ' '.join(WORDS[d] for d in digits)
    print(f"  '{src_str}' → '{tgt_str}'")
    print(f"   src_ids={src}  tgt_ids={tgt}")
    print()

# ── Small Seq2Seq Training ──
# In practice: define Encoder, Decoder, Seq2Seq as above, then:

class TinySeq2Seq(nn.Module):
    """Tiny Seq2Seq for the digit→words task."""
    def __init__(self, vocab_size=13, embed_dim=32, hidden_dim=64):
        super().__init__()
        self.enc_emb = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.dec_emb = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.encoder = nn.GRU(embed_dim, hidden_dim, batch_first=True)
        self.decoder = nn.GRUCell(embed_dim, hidden_dim)
        self.fc      = nn.Linear(hidden_dim, vocab_size)

    def forward(self, src, tgt, teacher_force=True):
        embedded_src = self.enc_emb(src)
        _, h = self.encoder(embedded_src)
        h = h.squeeze(0)  # (B, H)

        logits_all = []
        dec_input  = tgt[:, 0]  # SOS

        for t in range(1, tgt.size(1)):
            dec_emb = self.dec_emb(dec_input)
            h = self.decoder(dec_emb, h)
            logit = self.fc(h)
            logits_all.append(logit)
            dec_input = tgt[:, t] if teacher_force else logit.argmax(-1)

        return torch.stack(logits_all, dim=1)  # (B, T-1, V)

model     = TinySeq2Seq()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss(ignore_index=0)

# Quick training loop
for epoch in range(1, 201):
    # Generate batch
    batch_pairs = [make_pair(3) for _ in range(64)]
    max_src = max(len(p[1]) for p in batch_pairs)
    max_tgt = max(len(p[2]) for p in batch_pairs)

    src_batch = torch.zeros(64, max_src, dtype=torch.long)
    tgt_batch = torch.zeros(64, max_tgt, dtype=torch.long)
    for i, (_, src, tgt) in enumerate(batch_pairs):
        src_batch[i, :len(src)] = torch.tensor(src)
        tgt_batch[i, :len(tgt)] = torch.tensor(tgt)

    logits = model(src_batch, tgt_batch)
    loss   = criterion(logits.reshape(-1, 13), tgt_batch[:, 1:].reshape(-1))

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 50 == 0:
        print(f"Epoch {epoch:3d}: loss={loss.item():.4f}")

# Test the trained model
def translate(model, digit_str):
    src = torch.tensor([[SRC_VOCAB[d] for d in digit_str]])
    tgt = torch.tensor([[TGT_VOCAB['<SOS>']]])
    model.eval()
    with torch.no_grad():
        _, h = model.encoder(model.enc_emb(src))
        h = h.squeeze(0)
        result = []
        dec_input = torch.tensor([TGT_VOCAB['<SOS>']])
        for _ in range(10):
            dec_emb = model.dec_emb(dec_input)
            h = model.decoder(dec_emb, h)
            logit = model.fc(h)
            pred = logit.argmax(-1)
            if pred.item() == TGT_VOCAB['<EOS>']: break
            result.append(TGT_IDX_TO_WORD.get(pred.item(), '?'))
            dec_input = pred
    return ' '.join(result)

print("\nTranslation results:")
for test in ['123', '456', '789', '000']:
    print(f"  '{test}' → '{translate(model, test)}'")
# Expected: '123' → 'one two three', etc.