{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Word embeddings from scratch\n",
    "\n",
    "**GRA 4164 \u00b7 Practical day 2 \u00b7 Fall 2026**\n",
    "\n",
    "In lecture 6 we saw two routes to word embeddings: *counting* (co-occurrence \u2192 PPMI \u2192 SVD) and *predicting* (word2vec). This notebook walks the counting route end to end with nothing but numpy \u2014 on 223 State of the Union addresses (1790s\u2013today).\n",
    "\n",
    "By the end you will have built 50-dimensional word vectors that know that *war* goes with *enemy*, *gold* with *silver*, and *railroad* with *canal* \u2014 without labels, neural networks, or a GPU.\n",
    "\n",
    "**Setup:** download `sotu.zip` from the course site, unzip it so the speeches sit in `data/sotu/*.txt` next to this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from collections import Counter\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load the corpus\n",
    "\n",
    "One plain-text file per speech. We lowercase everything and keep only alphabetic tokens \u2014 the same preprocessing decisions as lecture 2, and remember: they are modelling decisions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = sorted(Path(\"data/sotu\").glob(\"*.txt\"))\n",
    "print(f\"{len(files)} speeches, e.g. {files[0].name} \u2026 {files[-1].name}\")\n",
    "\n",
    "tokens = []\n",
    "for p in files:\n",
    "    tokens.extend(re.findall(r\"[a-z]+\", p.read_text(errors=\"ignore\").lower()))\n",
    "print(f\"{len(tokens):,} tokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. The co-occurrence matrix\n",
    "\n",
    "Firth's idea made concrete: slide a window of \u00b14 words through the corpus and count, for each pair of words, how often they appear together. We keep the 2,000 most frequent words as our vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = [w for w, _ in Counter(tokens).most_common(2000)]\n",
    "idx = {w: i for i, w in enumerate(vocab)}\n",
    "V = len(vocab)\n",
    "\n",
    "Co = np.zeros((V, V), dtype=np.float32)\n",
    "window = 4\n",
    "ids = [idx.get(w, -1) for w in tokens]\n",
    "for i, wi in enumerate(ids):\n",
    "    if wi < 0:\n",
    "        continue\n",
    "    for j in range(max(0, i - window), min(len(ids), i + window + 1)):\n",
    "        if j != i and ids[j] >= 0:\n",
    "            Co[wi, ids[j]] += 1\n",
    "\n",
    "print(f\"co-occurrence matrix: {Co.shape}, {Co.sum():,.0f} pair counts\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. From counts to association: PPMI\n",
    "\n",
    "Raw counts are dominated by frequent words. Pointwise mutual information asks instead: *do these two words co-occur more than chance?*\n",
    "\n",
    "$$M_{ij} = \\log_2 \\frac{p_{ij}}{p_i \\, p_j}$$\n",
    "\n",
    "We keep only the positive part (PPMI) \u2014 negative associations are mostly noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total = Co.sum()\n",
    "p_i = Co.sum(axis=1, keepdims=True) / total\n",
    "p_j = Co.sum(axis=0, keepdims=True) / total\n",
    "\n",
    "with np.errstate(divide=\"ignore\"):\n",
    "    M = np.log2((Co / total) / (p_i * p_j))\n",
    "M[~np.isfinite(M)] = 0\n",
    "M = np.maximum(M, 0)\n",
    "\n",
    "print(f\"PPMI matrix: {(M > 0).mean():.1%} of entries positive\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Shrink it: truncated SVD\n",
    "\n",
    "The lecture-3 move on a new matrix: factorise, keep the top **K = 50** dimensions. Each row of $U_{1:K} S_{1:K}$ is a word's embedding. We normalise rows so dot products are cosine similarities (lecture 2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "U, S, _ = np.linalg.svd(M, full_matrices=False)\n",
    "\n",
    "K = 50\n",
    "emb = U[:, :K] * S[:K]\n",
    "emb = emb / np.linalg.norm(emb, axis=1, keepdims=True)\n",
    "print(f\"embeddings: {emb.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Ask the space questions\n",
    "\n",
    "Nearest neighbours by cosine similarity. Expect: *war \u2192 armed, attack, enemy*; *gold \u2192 silver, coin, specie*; *railroad \u2192 railway, canal*. The corpus spans two centuries \u2014 the vocabulary shows it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def neighbours(word, k=6):\n",
    "    sims = emb @ emb[idx[word]]\n",
    "    return [(vocab[i], float(sims[i])) for i in np.argsort(-sims)[1 : k + 1]]\n",
    "\n",
    "for w in [\"war\", \"tax\", \"china\", \"railroad\", \"gold\", \"peace\"]:\n",
    "    print(f\"{w:>10s} \u2192 \" + \", \".join(f\"{n} ({s:.2f})\" for n, s in neighbours(w)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Draw the space\n",
    "\n",
    "Project a few themed word groups to 2D (PCA on the selected words) \u2014 the same picture as the lecture-6 slide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "groups = {\n",
    "    \"war\":     [\"war\", \"army\", \"navy\", \"enemy\", \"military\", \"battle\"],\n",
    "    \"money\":   [\"tax\", \"revenue\", \"debt\", \"currency\", \"banks\", \"treasury\"],\n",
    "    \"law\":     [\"congress\", \"senate\", \"law\", \"constitution\", \"courts\"],\n",
    "    \"nations\": [\"mexico\", \"france\", \"spain\", \"britain\", \"china\", \"russia\"],\n",
    "}\n",
    "colors = {\"war\": \"tab:red\", \"money\": \"tab:olive\", \"law\": \"tab:blue\", \"nations\": \"tab:purple\"}\n",
    "\n",
    "words = [w for ws in groups.values() for w in ws]\n",
    "X = emb[[idx[w] for w in words]]\n",
    "Xc = X - X.mean(axis=0)\n",
    "_, _, Vt = np.linalg.svd(Xc, full_matrices=False)\n",
    "XY = Xc @ Vt[:2].T\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(9, 6))\n",
    "k = 0\n",
    "for g, ws in groups.items():\n",
    "    pts = XY[k : k + len(ws)]\n",
    "    ax.scatter(pts[:, 0], pts[:, 1], color=colors[g], label=g)\n",
    "    for (x, y), w in zip(pts, ws):\n",
    "        ax.annotate(w, (x, y), xytext=(5, 4), textcoords=\"offset points\", color=colors[g])\n",
    "    k += len(ws)\n",
    "ax.legend()\n",
    "ax.set_title(\"PPMI + SVD embeddings of the State of the Union corpus\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Exercises\n",
    "\n",
    "1. **Window size.** Rebuild with `window = 2` and `window = 10`. How do the neighbours of `war` change? (Lecture 6: small windows \u2192 syntactic, large \u2192 topical.)\n",
    "2. **Your own queries.** Find neighbours of words relevant to an assignment topic. Where does the space fail, and why might that be?\n",
    "3. **A cultural dimension (Kozlowski et al. 2019).** Build a direction `d = emb[idx[\"war\"]] - emb[idx[\"peace\"]]` and project other words onto it with `emb @ d`. Which words score highest and lowest? Try a `rich \u2212 poor` style pair of your own.\n",
    "4. **Time travel (harder).** Split the corpus into pre-1900 and post-1900 speeches, build separate embeddings, and compare the neighbours of `bank`, `union`, or `china` across the two."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}