{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "513561b0",
   "metadata": {},
   "source": [
    "# Sequence models: RNN, GRU, LSTM, attention and a small transformer\n",
    "\n",
    "Forecast a signal with RNN, GRU and LSTM, train a character-level language model and sample from it, play with embeddings, build attention from scratch and check it against PyTorch's fused kernel, train a tiny transformer that needs attention to succeed, compare greedy and beam search, and compute BLEU.\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/github/mahmoudyosrimahmoud13/inpace-library/blob/main/courses/dl-5-sequence-models/files/dl5_sequence_models_pytorch.ipynb) · part of the inPace course library. Slides © DeepLearning.AI (CC BY-SA 2.0); this notebook: inPace (CC BY-SA 4.0).\n",
    "\n",
    "**Run it top to bottom** (Runtime → Run all in Colab). Each part says what to look for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36cafc8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Works in Google Colab as it is (PyTorch is already installed there).\n",
    "# On your own computer:  pip install torch torchvision scikit-learn scipy matplotlib pandas\n",
    "# One framework per notebook: PyTorch and TensorFlow can clash when loaded into the same Python process.\n",
    "import numpy as np, matplotlib.pyplot as plt, torch\n",
    "from torch import nn\n",
    "torch.manual_seed(0); np.random.seed(0)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "print(\"PyTorch\", torch.__version__, \"| device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b96c743",
   "metadata": {},
   "source": [
    "## 1. RNN vs GRU vs LSTM on a sequence\n",
    "Predict the next value of a noisy signal from the last 40 values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecd01b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.arange(0, 400, 0.1)\n",
    "signal = np.sin(t) + 0.5 * np.sin(3.1 * t) + 0.1 * np.random.default_rng(0).normal(size=t.size)\n",
    "W = 40\n",
    "Xs = np.stack([signal[i:i + W] for i in range(len(signal) - W)])[:, :, None]\n",
    "ys = signal[W:]\n",
    "Xs, ys = torch.tensor(Xs, dtype=torch.float32), torch.tensor(ys, dtype=torch.float32)\n",
    "n = int(0.8 * len(Xs))\n",
    "class Forecaster(nn.Module):\n",
    "    def __init__(self, cell):\n",
    "        super().__init__(); self.rnn = cell(1, 32, batch_first=True); self.head = nn.Linear(32, 1)\n",
    "    def forward(self, x): return self.head(self.rnn(x)[0][:, -1]).squeeze(1)    # use the last time step\n",
    "for name, cell in ((\"RNN\", nn.RNN), (\"GRU\", nn.GRU), (\"LSTM\", nn.LSTM)):\n",
    "    torch.manual_seed(0); m = Forecaster(cell); opt = torch.optim.Adam(m.parameters(), lr=3e-3); losses = []\n",
    "    for epoch in range(3):\n",
    "        for i in torch.randperm(n).split(64):\n",
    "            opt.zero_grad(); loss = nn.functional.mse_loss(m(Xs[i]), ys[i]); loss.backward(); opt.step(); losses.append(loss.item())\n",
    "    with torch.no_grad(): test = nn.functional.mse_loss(m(Xs[n:]), ys[n:]).item()\n",
    "    plt.plot(np.convolve(losses, np.ones(20) / 20, \"valid\"), label=f\"{name}: test MSE {test:.4f}\")\n",
    "plt.yscale(\"log\"); plt.legend(); plt.title(\"training loss\"); plt.xlabel(\"mini-batch\"); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b35a9fac",
   "metadata": {},
   "source": [
    "## 2. A character-level language model\n",
    "Trained on a short text, it learns to predict the next character. Then we **sample** from it, one character at a time. Temperature < 1 makes it safer, > 1 more surprising."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e4f0a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = (\"slow and steady wins the race. the turtle kept walking while the hare slept. \"\n",
    "        \"every day a little progress, every lesson one stone on the trail. \") * 30\n",
    "chars = sorted(set(text)); stoi = {c: i for i, c in enumerate(chars)}\n",
    "data = torch.tensor([stoi[c] for c in text])\n",
    "class CharLM(nn.Module):\n",
    "    def __init__(self, v, e=32, h=128):\n",
    "        super().__init__(); self.emb = nn.Embedding(v, e); self.lstm = nn.LSTM(e, h, batch_first=True); self.out = nn.Linear(h, v)\n",
    "    def forward(self, x, state=None):\n",
    "        y, state = self.lstm(self.emb(x), state); return self.out(y), state\n",
    "torch.manual_seed(0); lm = CharLM(len(chars)); opt = torch.optim.AdamW(lm.parameters(), lr=3e-3)\n",
    "for step in range(600):\n",
    "    i = torch.randint(0, len(data) - 65, (32,))\n",
    "    x = torch.stack([data[j:j + 64] for j in i]); y = torch.stack([data[j + 1:j + 65] for j in i])\n",
    "    logits, _ = lm(x); loss = nn.functional.cross_entropy(logits.reshape(-1, len(chars)), y.reshape(-1))\n",
    "    opt.zero_grad(); loss.backward(); opt.step()\n",
    "print(\"final loss\", round(loss.item(), 3))\n",
    "def sample(start=\"the \", n=120, temperature=0.8):\n",
    "    lm.eval(); out = start; x = torch.tensor([[stoi[c] for c in start]]); state = None\n",
    "    with torch.no_grad():\n",
    "        for _ in range(n):\n",
    "            logits, state = lm(x, state)\n",
    "            c = torch.multinomial(torch.softmax(logits[0, -1] / temperature, -1), 1).item()\n",
    "            out += chars[c]; x = torch.tensor([[c]])\n",
    "    return out\n",
    "for temp in (0.5, 1.0, 1.5): print(f\"T={temp}: {sample(temperature=temp)!r}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da1832c8",
   "metadata": {},
   "source": [
    "## 3. Embeddings and cosine similarity\n",
    "`nn.Embedding` is just a lookup table (the slides' embedding matrix E, one row per word). A tiny hand-made example shows the famous analogy arithmetic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159e919e",
   "metadata": {},
   "outputs": [],
   "source": [
    "words = [\"man\", \"woman\", \"king\", \"queen\", \"apple\", \"orange\"]\n",
    "#            gender royal  food\n",
    "vec = torch.tensor([[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [-0.95, 0.9, 0.0], [0.97, 0.95, 0.0], [0.0, 0.0, 1.0], [0.05, 0.0, 0.95]])\n",
    "emb = nn.Embedding.from_pretrained(vec)\n",
    "e = lambda w: emb(torch.tensor(words.index(w)))\n",
    "target = e(\"king\") - e(\"man\") + e(\"woman\")\n",
    "sims = {w: nn.functional.cosine_similarity(target, e(w), dim=0).item() for w in words if w not in (\"king\", \"man\", \"woman\")}\n",
    "print(\"king - man + woman is closest to:\", max(sims, key=sims.get), {k: round(v, 2) for k, v in sims.items()})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0074d8a",
   "metadata": {},
   "source": [
    "## 4. Attention from scratch, and PyTorch's fused version\n",
    "`Attention(Q, K, V) = softmax(QKᵀ / √d) V`. PyTorch's `scaled_dot_product_attention` does the same in one fast call (with FlashAttention-style kernels on GPUs). A **causal mask** stops a word from looking at later words (what GPT-style models use)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "181fad2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def attention(Q, K, V, causal=False):\n",
    "    scores = Q @ K.transpose(-2, -1) / Q.shape[-1] ** 0.5\n",
    "    if causal:\n",
    "        L = scores.shape[-1]\n",
    "        scores = scores.masked_fill(torch.triu(torch.ones(L, L, dtype=torch.bool), 1), float(\"-inf\"))\n",
    "    return torch.softmax(scores, -1) @ V\n",
    "Q, K, V = torch.randn(2, 5, 16), torch.randn(2, 5, 16), torch.randn(2, 5, 16)\n",
    "F = nn.functional\n",
    "print(\"matches PyTorch:\", torch.allclose(attention(Q, K, V), F.scaled_dot_product_attention(Q, K, V), atol=1e-5))\n",
    "print(\"causal matches:\", torch.allclose(attention(Q, K, V, True), F.scaled_dot_product_attention(Q, K, V, is_causal=True), atol=1e-5))\n",
    "mha = nn.MultiheadAttention(embed_dim=16, num_heads=4, batch_first=True)\n",
    "out, w = mha(Q, K, V)\n",
    "print(\"multi-head output\", tuple(out.shape), \"| attention weights\", tuple(w.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09fd6776",
   "metadata": {},
   "source": [
    "## 5. A tiny transformer that needs attention\n",
    "Task: for every position, output **how many times its token appears in the whole sequence**. A model that looks at one position at a time can't do it; attention can, because every position can look at every other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "444e1781",
   "metadata": {},
   "outputs": [],
   "source": [
    "V_, L_ = 6, 12\n",
    "def batch(n):\n",
    "    x = torch.randint(0, V_, (n, L_))\n",
    "    y = (x[:, :, None] == x[:, None, :]).sum(-1) - 1           # count of the same token, 1..L -> 0..L-1\n",
    "    return x, y\n",
    "class TinyTransformer(nn.Module):\n",
    "    def __init__(self, d=64):\n",
    "        super().__init__()\n",
    "        self.tok = nn.Embedding(V_, d); self.pos = nn.Embedding(L_, d)\n",
    "        layer = nn.TransformerEncoderLayer(d, nhead=4, dim_feedforward=128, batch_first=True, dropout=0.0)\n",
    "        self.enc = nn.TransformerEncoder(layer, num_layers=2); self.out = nn.Linear(d, L_)\n",
    "    def forward(self, x):\n",
    "        return self.out(self.enc(self.tok(x) + self.pos(torch.arange(x.shape[1]))))\n",
    "class NoAttention(nn.Module):                                  # same size, but each position alone\n",
    "    def __init__(self, d=64):\n",
    "        super().__init__(); self.net = nn.Sequential(nn.Embedding(V_, d), nn.Linear(d, 128), nn.ReLU(), nn.Linear(128, L_))\n",
    "    def forward(self, x): return self.net(x)\n",
    "for name, m in ((\"no attention\", NoAttention()), (\"transformer\", TinyTransformer())):\n",
    "    torch.manual_seed(0); opt = torch.optim.AdamW(m.parameters(), lr=1e-3)\n",
    "    for step in range(1500):\n",
    "        x, y = batch(64); loss = F.cross_entropy(m(x).reshape(-1, L_), y.reshape(-1))\n",
    "        opt.zero_grad(); loss.backward(); opt.step()\n",
    "    x, y = batch(1000)\n",
    "    with torch.no_grad(): print(f\"{name:<13} accuracy {(m(x).argmax(-1) == y).float().mean():.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c4d5bd",
   "metadata": {},
   "source": [
    "## 7. Greedy search vs beam search\n",
    "A toy 'translation model' as a table of next-word probabilities. Greedy takes the best word each time; beam search keeps the B best partial sentences and finds a more likely whole sentence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4c2cfcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "P = {\"<s>\": {\"jane\": 0.5, \"in\": 0.4, \"the\": 0.1}, \"jane\": {\"is\": 0.5, \"visits\": 0.5},\n",
    "     \"in\": {\"september\": 1.0}, \"is\": {\"going\": 0.6, \"visiting\": 0.4}, \"visits\": {\"africa\": 0.9, \"</s>\": 0.1},\n",
    "     \"going\": {\"to\": 0.55, \"</s>\": 0.45}, \"to\": {\"be\": 0.5, \"</s>\": 0.5}, \"be\": {\"</s>\": 1.0}, \"visiting\": {\"africa\": 0.9, \"</s>\": 0.1},\n",
    "     \"africa\": {\"</s>\": 1.0}, \"september\": {\"</s>\": 1.0}, \"the\": {\"</s>\": 1.0}}\n",
    "def greedy():\n",
    "    s, lp = [\"<s>\"], 0.0\n",
    "    while s[-1] != \"</s>\":\n",
    "        w, p = max(P[s[-1]].items(), key=lambda kv: kv[1]); s.append(w); lp += math.log(p)\n",
    "    return s[1:-1], lp\n",
    "def beam(B=3):\n",
    "    beams = [([\"<s>\"], 0.0)]\n",
    "    while any(b[0][-1] != \"</s>\" for b in beams):\n",
    "        cand = []\n",
    "        for s, lp in beams:\n",
    "            if s[-1] == \"</s>\": cand.append((s, lp)); continue\n",
    "            cand += [(s + [w], lp + math.log(p)) for w, p in P[s[-1]].items()]\n",
    "        beams = sorted(cand, key=lambda b: b[1] / len(b[0]) ** 0.7, reverse=True)[:B]   # length-normalized\n",
    "    return beams[0][0][1:-1], beams[0][1]\n",
    "for name, (s, lp) in ((\"greedy\", greedy()), (\"beam B=3\", beam())):\n",
    "    print(f\"{name:<9} {' '.join(s):<28} probability {math.exp(lp):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68883967",
   "metadata": {},
   "source": [
    "## 8. BLEU score\n",
    "Modified n-gram precision, combined over 1- to 4-grams, with a brevity penalty for short outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1aca372",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "def bleu(candidate, references, N=4):\n",
    "    c = candidate.split(); refs = [r.split() for r in references]; logs = []\n",
    "    for n in range(1, N + 1):\n",
    "        grams = Counter(tuple(c[i:i + n]) for i in range(len(c) - n + 1))\n",
    "        clip = sum(min(cnt, max(Counter(tuple(r[i:i + n]) for i in range(len(r) - n + 1))[g] for r in refs)) for g, cnt in grams.items())\n",
    "        logs.append(math.log(max(clip, 1e-9) / max(sum(grams.values()), 1)))\n",
    "    ref_len = min((abs(len(r) - len(c)), len(r)) for r in refs)[1]\n",
    "    bp = 1.0 if len(c) > ref_len else math.exp(1 - ref_len / max(len(c), 1))\n",
    "    return bp * math.exp(sum(logs) / N)\n",
    "refs = [\"the cat is on the mat\", \"there is a cat on the mat\"]\n",
    "for cand in (\"the cat is on the mat\", \"the cat on the mat\", \"the the the the the the\"):\n",
    "    print(f\"{cand:<26} BLEU {bleu(cand, refs, N=2):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "290007ca",
   "metadata": {},
   "source": [
    "## 9. Using a pre-trained transformer (Hugging Face)\n",
    "Today you rarely train language models from zero: you start from a pre-trained one. Runs in Colab (downloads a small model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4a18ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from transformers import pipeline\n",
    "    clf = pipeline(\"sentiment-analysis\", model=\"distilbert/distilbert-base-uncased-finetuned-sst-2-english\")\n",
    "    print(clf([\"The dessert is excellent.\", \"Completely lacking in good taste, good service and good ambience.\"]))\n",
    "except Exception as e:\n",
    "    print(\"Run this in Colab (needs `transformers` and internet):\", type(e).__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49250e12",
   "metadata": {},
   "source": [
    "## Try this\n",
    "1. In section 1, use 200 time steps instead of 40. Which cell suffers most?\n",
    "2. In section 2, sample with temperature 0.1 and 3.0. Describe what happens.\n",
    "3. In section 5, give the transformer 1 layer and 1 head. Does it still solve the task?\n",
    "4. In section 7, try beam width 1 and 5. Which equals greedy search?"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}