{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d45605ec-3f07-40d8-ab31-10a38ce83b79",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/7.lm/HW6_Perplexity.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "397af444-94c2-477b-92b8-4467b5ea80af",
   "metadata": {},
   "source": [
    "# HW6: Perplexity\n",
    "\n",
    "In this homework, you will implement a function to calculate the perplexity of the n-gram language models we covered in lab, and experiment with different sequences to better understand both n-gram LMs as well as the perplexity metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52b3a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "from collections import Counter\n",
    "\n",
    "import nltk\n",
    "import numpy as np\n",
    "from nltk import sent_tokenize, word_tokenize\n",
    "\n",
    "nltk.download(\"punkt\")\n",
    "nltk.download(\"punkt_tab\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33eb376a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_file(filename):\n",
    "    sequences = []\n",
    "    with open(filename) as file:\n",
    "        data = file.read()\n",
    "        sents = sent_tokenize(data)\n",
    "        for sent in sents:\n",
    "            tokens = word_tokenize(sent)\n",
    "            sequences.append(tokens)\n",
    "    return sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a66e06bf-948f-4323-9af8-5864b466b033",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget --no-check-certificate https://raw.githubusercontent.com/dbamman/anlp25/refs/heads/main/data/1342_pride_and_prejudice.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "feeb46f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read data from file and tokenize them into sequences comprised of tokens.\n",
    "\n",
    "# Pride and Prejudice (Jane Austen)\n",
    "sequences = read_file(\"1342_pride_and_prejudice.txt\")\n",
    "\n",
    "max_sequences = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "214dbd04",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NgramModel():\n",
    "\n",
    "    def __init__(self, sequences, order):\n",
    "\n",
    "        # For this exercise we're going to encode the LM as a sparse dictionary (trading less storage for more compute)\n",
    "        # We'll store the LM as a dictionary with the conditioning context as keys; each value is a\n",
    "        # Counter object that keeps track of the number of times we see a word following that context.\n",
    "        self.counts = {}\n",
    "\n",
    "        # Markov order (order 1 = conditioning on previous 1 word; order 2 = previous 2 words, etc.)\n",
    "        self.order = order\n",
    "\n",
    "        vocab = {\"[END]\": 0}\n",
    "\n",
    "        for s_idx, tokens in enumerate(sequences):\n",
    "            # We'll add [START] and [END] tokens to encode the beginning/end of sentences\n",
    "            tokens = [\"[START]\"] * order + tokens + [\"[END]\"]\n",
    "\n",
    "            if s_idx == 0:\n",
    "                print(tokens)\n",
    "\n",
    "            for i in range(order, len(tokens)):\n",
    "                context = \" \".join(tokens[i - order:i])\n",
    "                word = tokens[i]\n",
    "\n",
    "                if word not in vocab:\n",
    "                    vocab[word] = len(vocab)\n",
    "\n",
    "                # For just the first sentence, print the conditioning context + word\n",
    "                if s_idx == 0:\n",
    "                    print(\"Context: %s Next: %s\" % (context.ljust(50), word))\n",
    "\n",
    "                if context not in self.counts:\n",
    "                    self.counts[context] = Counter()\n",
    "                self.counts[context][word] += 1\n",
    "\n",
    "\n",
    "\n",
    "    def sample(self, context):\n",
    "        total = sum(self.counts[context].values())\n",
    "\n",
    "        dist = []\n",
    "        vocab = []\n",
    "\n",
    "        # Create a probability distribution for each conditioning context, over the vocab that we've observed it with.\n",
    "        for idx, word in enumerate(self.counts[context]):\n",
    "            prob = self.counts[context][word]/total\n",
    "            dist.append(prob)\n",
    "            vocab.append(word)\n",
    "\n",
    "        index = np.argmax(np.random.multinomial(1, pvals=dist))\n",
    "        return vocab[index]\n",
    "\n",
    "    def generate_sequence(self, keep_ends=True):\n",
    "        generated = [\"[START]\"] * (self.order)\n",
    "        word = None\n",
    "        while word != \"[END]\":\n",
    "            context = ' '.join(generated[-self.order:] if self.order > 0 else \"\")\n",
    "            word = self.sample(context)\n",
    "            generated.append(word)\n",
    "        if not keep_ends:\n",
    "            generated = generated[self.order:-1]\n",
    "        return \" \".join(generated)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0384b034",
   "metadata": {},
   "source": [
    "Let's create some language models of different orders from *Pride and Prejudice*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4502623",
   "metadata": {},
   "outputs": [],
   "source": [
    "ngram0 = NgramModel(sequences[:max_sequences], order=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81a51b9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ngram1 = NgramModel(sequences[:max_sequences], order=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35ee1ead",
   "metadata": {},
   "outputs": [],
   "source": [
    "ngram2 = NgramModel(sequences[:max_sequences], order=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c55ac8d",
   "metadata": {},
   "source": [
    "**Q1.** Create a `perplexity` function that can takes two arguments: a.) a model of *any* ngram order (from the class above); and b.) a sequence to calculate perplexity for.  You'll recall from class that perplexity under a particular language model for sequence $w$ is given by the following equation:\n",
    "\n",
    "$$\n",
    "\\textrm{perplexity}_{model}(w) = \\exp\\left(-{1 \\over N} \\sum_{i=1}^N \\log P_{model}(w_i) \\right)\n",
    "$$\n",
    "\n",
    "$P_{model}(w_i)$ calculates the probability of token $w_i$ using whatever assumptions that model makes -- for a bigram model (order 1), this is $P(w_i \\mid w_{i-1})$, for a trigram model (order 2), this is $P(w_i \\mid w_{i-2}, w_{i-1})$, etc.  Two things to note:\n",
    "\n",
    "* When calculating the probability of the first word(s), be sure to get the conditioning context right.  The conditioning context for the first word in a trigram model, for example, is $P(w_i \\mid$ [START] [START]$)$.\n",
    "* Perplexity is only calculated for the words in the actual sequence.  We don't include the $P$([START]) or $P$([END]) in the perplexity calculuation.\n",
    "\n",
    "\n",
    "*Hint*: when working on this function, you might want to debug by printing out the probabilities of each $w_i$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d32301c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import log\n",
    "\n",
    "def perplexity(model, tokens):\n",
    "    # YOUR CODE HERE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea8d0a8a",
   "metadata": {},
   "source": [
    "**Q2**. Execute that perplexity function for the following language models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a4b0e8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "perplexity(ngram0, word_tokenize(\"She was a great friend of Mr. Bingley.\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd2d4b64",
   "metadata": {},
   "outputs": [],
   "source": [
    "perplexity(ngram1, word_tokenize(\"She was a great friend of Mr. Bingley.\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05be627d",
   "metadata": {},
   "outputs": [],
   "source": [
    "perplexity(ngram2, word_tokenize(\"She was a great friend of Mr. Bingley.\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a331078",
   "metadata": {},
   "source": [
    "**Q3.** What is the perplexity of \"She was a really great friend of Mr. Bingley.\" in the trigram language model trained above?\n",
    "\n",
    "Explain in 100 words what behavior is expected (and correct) given how an n-gram language model works and the data we are training it on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a4f81c",
   "metadata": {},
   "outputs": [],
   "source": [
    "perplexity(ngram2, word_tokenize(\"She was a really great friend of Mr. Bingley.\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19d3de75-a37c-4983-8f40-a8306142caca",
   "metadata": {},
   "source": [
    "**Q4.** What 1-token sequence yields the lowest perplexity for the 0-order n-gram model? Why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7143c48-5c08-4ada-af35-8488b0de3ffd",
   "metadata": {},
   "outputs": [],
   "source": [
    "perplexity(ngram0, word_tokenize(\"YOUR_ANSWER_HERE\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4e059ed-8d4f-4698-b0c2-c42231a2105a",
   "metadata": {},
   "source": [
    "**Q5.** Write a function to find the n-token sequence with the lowest perplexity given an n-gram model. Explain why it should work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffb4d492-8c34-47f9-bf3f-b03ca2dfcdd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def min_ppl_sequence(model, n=5):\n",
    "    \"\"\"Return an n-token long string that yields the lowest possible perplexity under the provided model.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3facf28b-323e-4f21-b736-24d707d10f6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_ppl_sequence(ngram1, n=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c13d61-547c-427f-9f0e-bee7c8d27f59",
   "metadata": {},
   "outputs": [],
   "source": [
    "perplexity(ngram1, min_ppl_sequence(ngram1, n=3))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}