{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# tf-idf and PMI, by hand\n",
    "\n",
    "The weighting arithmetic the chapters print, on the four document matrix.\n",
    "\n",
    "From chapter 4, [Term Weighting and Similarity](https://nlp.jcrlabz.com/book/termweight/), of the course notes.\n",
    "\n",
    "Source: `book/code/worked_examples/weighting.py`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What you need\n",
    "\n",
    "Nothing. This example uses only the Python standard library."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Nothing to install: this example uses only the Python standard library.\n",
    "import sys; print(sys.version.split()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The script\n",
    "\n",
    "Everything the example defines, in one cell. Run it and the definitions\n",
    "are live in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "\"\"\"tf-idf and PMI, worked out as Chapters 4 and 7 work them out.\n",
    "\n",
    "Two ideas, both of which come down to dividing by what you expected.\n",
    "\n",
    "    tf-idf  divides a word's count in a document by how many documents hold it\n",
    "    PMI     divides a pair's joint count by what independence would predict\n",
    "\n",
    "    python3 weighting.py tfidf     # the moon / the worked example\n",
    "    python3 weighting.py corpus    # the four-document matrix and cosines\n",
    "    python3 weighting.py pmi       # the ice cream worked example\n",
    "    python3 weighting.py all       # all three\n",
    "\n",
    "Run `pmi` and watch the word \"the\" fail to score, which is the whole point.\n",
    "Run `corpus` and watch two documents that share only \"the\" come out at a\n",
    "cosine similarity of exactly zero.\n",
    "\n",
    "Install:  nothing, the Python standard library is enough\n",
    "\"\"\"\n",
    "\n",
    "import argparse\n",
    "import itertools\n",
    "import math\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "# --------------------------------------------------------------- tf-idf\n",
    "\n",
    "def idf(n_docs, doc_freq, base=10):\n",
    "    \"\"\"How informative a term is across the collection.\n",
    "\n",
    "    A term in every document scores 0, because N/df = 1 and log 1 = 0. It\n",
    "    cannot tell any two documents apart, so its weight is annihilated.\n",
    "    \"\"\"\n",
    "    return math.log(n_docs / doc_freq, base)\n",
    "\n",
    "\n",
    "def tf_idf(count_in_doc, doc_length, n_docs, doc_freq, base=10):\n",
    "    \"\"\"Length-normalised term frequency times inverse document frequency.\"\"\"\n",
    "    tf = count_in_doc / doc_length\n",
    "    return tf * idf(n_docs, doc_freq, base)\n",
    "\n",
    "\n",
    "def demo_tfidf():\n",
    "    N = 100_000\n",
    "    print(f\"\\ntf-idf   collection of N = {N:,} documents\")\n",
    "    print(\"         one document, 427 tokens long\\n\")\n",
    "    rows = [\n",
    "        # word,      count in doc, document frequency\n",
    "        (\"moon\",     20,   100),\n",
    "        (\"the\",      40,   N),\n",
    "        (\"telescope\", 3,   4_000),\n",
    "        (\"crater\",    1,   60),\n",
    "    ]\n",
    "    print(f\"  {'word':<12}{'count':>7}{'df':>9}{'tf':>9}{'idf':>7}{'tf-idf':>9}\")\n",
    "    print(\"  \" + \"-\" * 53)\n",
    "    for word, count, df in rows:\n",
    "        tf = count / 427\n",
    "        i = idf(N, df)\n",
    "        print(f\"  {word:<12}{count:>7}{df:>9,}{tf:>9.4f}{i:>7.2f}\"\n",
    "              f\"{tf * i:>9.4f}\")\n",
    "\n",
    "    print(\"\\n  Read the last two rows against the first two.\")\n",
    "    print(\"  'the' occurs twice as often as 'moon' and scores ZERO, because\")\n",
    "    print(\"  it is in every document. 'crater' occurs once and outscores it.\")\n",
    "    print(\"  Raw frequency is the wrong guide; tf-idf is a machine for\")\n",
    "    print(\"  discounting the Zipfian head.\")\n",
    "\n",
    "\n",
    "# -------------------------------------------------- the four-document corpus\n",
    "\n",
    "DOCS = {\n",
    "    \"d1\": \"the moon orbits the earth\",\n",
    "    \"d2\": \"the earth orbits the sun\",\n",
    "    \"d3\": \"the telescope shows the moon\",\n",
    "    \"d4\": \"the sun is a star\",\n",
    "}\n",
    "# Listed in this order so the printed matrix matches the book's table.\n",
    "VOCAB = [\"the\", \"moon\", \"orbits\", \"earth\", \"sun\",\n",
    "         \"telescope\", \"shows\", \"is\", \"a\", \"star\"]\n",
    "\n",
    "\n",
    "def cosine(u, v):\n",
    "    \"\"\"Angle between two vectors, ignoring their lengths.\"\"\"\n",
    "    nu = math.sqrt(sum(x * x for x in u))\n",
    "    nv = math.sqrt(sum(x * x for x in v))\n",
    "    if nu == 0 or nv == 0:\n",
    "        return 0.0\n",
    "    return sum(a * b for a, b in zip(u, v)) / (nu * nv)\n",
    "\n",
    "\n",
    "def demo_corpus():\n",
    "    names = list(DOCS)\n",
    "    counts = {d: Counter(DOCS[d].split()) for d in names}\n",
    "    n_docs = len(names)\n",
    "    df = {w: sum(1 for d in names if counts[d][w] > 0) for w in VOCAB}\n",
    "    idf = {w: math.log10(n_docs / df[w]) for w in VOCAB}\n",
    "\n",
    "    print(f\"\\nA four-document collection, N = {n_docs}\\n\")\n",
    "    for d in names:\n",
    "        print(f\"  {d}: {DOCS[d]}\")\n",
    "\n",
    "    print(\"\\nStage 1. Count. Rows are terms, columns are documents.\\n\")\n",
    "    head = \"\".join(f\"{d:>6}\" for d in names)\n",
    "    print(f\"  {'term':<11}{head}\")\n",
    "    print(\"  \" + \"-\" * (11 + 6 * n_docs))\n",
    "    for w in VOCAB:\n",
    "        print(f\"  {w:<11}\" + \"\".join(f\"{counts[d][w]:>6}\" for d in names))\n",
    "\n",
    "    print(\"\\nStage 2. Count the documents, not the occurrences.\")\n",
    "    print(\"         df is how many documents hold the term at all.\")\n",
    "    print(\"         idf = log10(N / df).\\n\")\n",
    "    print(f\"  {'term':<11}{'df':>4}{'N/df':>8}{'idf':>8}\")\n",
    "    print(\"  \" + \"-\" * 31)\n",
    "    for w in VOCAB:\n",
    "        print(f\"  {w:<11}{df[w]:>4}{n_docs/df[w]:>8.2f}{idf[w]:>8.3f}\")\n",
    "    print(\"\\n  'the' is in every document, so N/df = 1 and idf = 0.\")\n",
    "    print(\"  Its weight is annihilated before any similarity is computed.\")\n",
    "\n",
    "    print(\"\\nStage 3. Multiply, cell by cell. tf-idf = tf x idf.\\n\")\n",
    "    weights = {d: [counts[d][w] * idf[w] for w in VOCAB] for d in names}\n",
    "    print(f\"  {'term':<11}\" + \"\".join(f\"{d:>8}\" for d in names))\n",
    "    print(\"  \" + \"-\" * (11 + 8 * n_docs))\n",
    "    for i, w in enumerate(VOCAB):\n",
    "        print(f\"  {w:<11}\" + \"\".join(f\"{weights[d][i]:>8.3f}\" for d in names))\n",
    "    print(\"  \" + \"-\" * (11 + 8 * n_docs))\n",
    "    print(f\"  {'norm':<11}\" + \"\".join(\n",
    "        f\"{math.sqrt(sum(x*x for x in weights[d])):>8.4f}\" for d in names))\n",
    "    print(\"\\n  The whole first row is zero. Two documents that share only\")\n",
    "    print(\"  'the' now share nothing at all.\")\n",
    "\n",
    "    print(\"\\nStage 4. Compare, as an angle.\\n\")\n",
    "    raw = {d: [counts[d][w] for w in VOCAB] for d in names}\n",
    "    print(f\"  {'pair':<10}{'cos(raw tf)':>13}{'cos(tf-idf)':>14}\")\n",
    "    print(\"  \" + \"-\" * 37)\n",
    "    for a, b in itertools.combinations(names, 2):\n",
    "        print(f\"  {a + '-' + b:<10}{cosine(raw[a], raw[b]):>13.4f}\"\n",
    "              f\"{cosine(weights[a], weights[b]):>14.4f}\")\n",
    "\n",
    "    print(\"\\n  Read d2-d3. Raw counts call them 0.57 similar. They share\")\n",
    "    print(\"  exactly one word, and that word is 'the'. Under tf-idf their\")\n",
    "    print(\"  cosine is 0.0000, which is the correct answer.\")\n",
    "    print(\"  Read d1-d2 against d1-d3. Raw counts separate them by a factor\")\n",
    "    print(\"  of 1.2. tf-idf separates them by a factor of 3.5. The ranking\")\n",
    "    print(\"  was already right; the weighting made it decisive.\")\n",
    "\n",
    "    print(\"\\nStage 5. Rank documents for a query.\")\n",
    "    print(\"         score(q, d) = sum over t in q of tf-idf(t, d).\\n\")\n",
    "    for query in (\"the moon\", \"moon orbits\"):\n",
    "        qt = query.split()\n",
    "        scored = sorted(\n",
    "            ((sum(counts[d][w] * idf[w] for w in qt),\n",
    "              sum(counts[d][w] for w in qt), d) for d in names),\n",
    "            key=lambda r: -r[0])\n",
    "        print(f\"  query \\\"{query}\\\"\")\n",
    "        print(f\"    {'doc':<5}{'raw count':>11}{'tf-idf score':>14}\")\n",
    "        for score, rawc, d in scored:\n",
    "            print(f\"    {d:<5}{rawc:>11}{score:>14.4f}\")\n",
    "    print(\"\\n  On \\\"the moon\\\", raw counts tie d1 and d3 at 3 with d2 close\")\n",
    "    print(\"  behind at 2. tf-idf keeps the tie between d1 and d3, which is\")\n",
    "    print(\"  right, and sends d2 and d4 to exactly zero, which is also\")\n",
    "    print(\"  right. Neither of them mentions the moon.\")\n",
    "\n",
    "    print(\"\\n  One honest caveat. In this collection 'a' has df = 1, so it\")\n",
    "    print(\"  scores the highest idf there is. idf is a statistic of the\")\n",
    "    print(\"  collection, not a judgement about language. With N = 4 it is\")\n",
    "    print(\"  measuring almost nothing. Give it 100,000 documents and 'a'\")\n",
    "    print(\"  falls to zero alongside 'the'.\")\n",
    "\n",
    "\n",
    "# ------------------------------------------------------------------ PMI\n",
    "\n",
    "def pmi(joint_count, count_w, count_c, n_tokens):\n",
    "    \"\"\"log2 of (what we saw) over (what independence predicts).\n",
    "\n",
    "    0 means exactly independent. Positive means the pair sticks together\n",
    "    more than chance. Negative means less, and in a sparse matrix negative\n",
    "    values are mostly noise, which is why PPMI clips them to 0.\n",
    "    \"\"\"\n",
    "    return math.log2((joint_count * n_tokens) / (count_w * count_c))\n",
    "\n",
    "\n",
    "def ppmi(*args):\n",
    "    return max(pmi(*args), 0.0)\n",
    "\n",
    "\n",
    "def demo_pmi():\n",
    "    N = 1_000_000\n",
    "    print(f\"\\nPMI      corpus of N = {N:,} tokens\\n\")\n",
    "    rows = [\n",
    "        # w,      c,        joint, count_w,  count_c\n",
    "        (\"ice\",   \"cream\",    500,   2_000,    2_000),   # strong\n",
    "        (\"ice\",   \"the\",      800,   2_000,   50_000),   # frequent, not strong\n",
    "        (\"data\",  \"science\",   10,   1_000,    1_000),   # modest\n",
    "        (\"cold\",  \"cream\",      4,   2_000,    2_000),   # exactly independent\n",
    "        (\"cream\", \"asphalt\",    1,   2_000,    2_000),   # below chance\n",
    "    ]\n",
    "    print(f\"  {'pair':<18}{'joint':>7}{'expected':>10}{'PMI':>8}{'PPMI':>7}\")\n",
    "    print(\"  \" + \"-\" * 50)\n",
    "    for w, c, joint, cw, cc in rows:\n",
    "        expected = cw * cc / N          # what independence predicts\n",
    "        print(f\"  {w + ' + ' + c:<18}{joint:>7}{expected:>10.1f}\"\n",
    "              f\"{pmi(joint, cw, cc, N):>8.2f}{ppmi(joint, cw, cc, N):>7.2f}\")\n",
    "\n",
    "    print(\"\\n  Row 2 is the one to study. 'ice' sits next to 'the' 800 times,\")\n",
    "    print(\"  far more often than it sits next to 'cream'. Raw counts would\")\n",
    "    print(\"  call that the stronger association.\")\n",
    "    print(\"  PMI does not. 'the' is common on its own, so the expected count\")\n",
    "    print(\"  is large too, and the ratio collapses from 800 down to 3 bits.\")\n",
    "    print(\"\\n  Row 4 is exactly independent. Joint equals expected, the ratio\")\n",
    "    print(\"  is 1, and log2(1) = 0. That is the reading to memorise: PMI of\")\n",
    "    print(\"  zero means the pair tells you nothing.\")\n",
    "    print(\"\\n  Row 5 is below chance, so PMI goes negative. PPMI clips it to 0.\")\n",
    "    print(\"  Seeing a pair once is no evidence that it is avoided; it is much\")\n",
    "    print(\"  more likely that the corpus is simply too small to say.\")\n",
    "\n",
    "\n",
    "def main():\n",
    "    ap = argparse.ArgumentParser(description=__doc__)\n",
    "    ap.add_argument(\"which\", nargs=\"?\", default=\"all\",\n",
    "                    choices=[\"tfidf\", \"corpus\", \"pmi\", \"all\"])\n",
    "    a = ap.parse_args()\n",
    "    if a.which in (\"tfidf\", \"all\"):\n",
    "        demo_tfidf()\n",
    "    if a.which in (\"corpus\", \"all\"):\n",
    "        demo_corpus()\n",
    "    if a.which in (\"pmi\", \"all\"):\n",
    "        demo_pmi()\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running it\n",
    "\n",
    "The script is a command line tool, so each run below sets the arguments\n",
    "and calls `main()`. The heading shows the equivalent shell command."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "def run(*args):\n",
    "    sys.argv = [\"weighting.py\", *args]\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 weighting.py`\n",
    "\n",
    "Every section at once."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 weighting.py pmi`\n",
    "\n",
    "ice sits beside the 800 times and beside cream 500 times. Raw counts call the the stronger association. PMI does not, because it divides by what independence predicts."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"pmi\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `python3 weighting.py corpus`\n",
    "\n",
    "Two documents that share only the word the come out at a cosine of exactly zero."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "run(\"corpus\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Read the chapter this comes from: [Term Weighting and Similarity](https://nlp.jcrlabz.com/book/termweight/)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "weighting.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}