{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d423c6e",
   "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/ExploreLM.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf7e7a1-39e7-4dc6-9450-769112d4c6b1",
   "metadata": {},
   "source": [
    "# Language modeling\n",
    "\n",
    "In this notebook, we will construct an n-gram language model, and use it to generate sequences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52b3a60",
   "metadata": {
    "id": "b52b3a60"
   },
   "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": {
    "id": "33eb376a"
   },
   "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": "_4P6ewmTfAo0",
   "metadata": {
    "id": "_4P6ewmTfAo0"
   },
   "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": {
    "id": "feeb46f4"
   },
   "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": {
    "id": "214dbd04"
   },
   "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": "code",
   "execution_count": null,
   "id": "81a51b9e",
   "metadata": {
    "id": "81a51b9e"
   },
   "outputs": [],
   "source": [
    "ngram1 = NgramModel(sequences[:max_sequences], order=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79935c2",
   "metadata": {
    "id": "c79935c2"
   },
   "outputs": [],
   "source": [
    "print(ngram1.generate_sequence())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad8a2c8d",
   "metadata": {
    "id": "ad8a2c8d"
   },
   "outputs": [],
   "source": [
    "ngram0 = NgramModel(sequences[:max_sequences], order=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6982433e",
   "metadata": {
    "id": "6982433e"
   },
   "outputs": [],
   "source": [
    "ngram0.generate_sequence()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35ee1ead",
   "metadata": {
    "id": "35ee1ead"
   },
   "outputs": [],
   "source": [
    "ngram2 = NgramModel(sequences[:max_sequences], order=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24378b79",
   "metadata": {
    "id": "24378b79"
   },
   "outputs": [],
   "source": [
    "ngram2.generate_sequence(keep_ends=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49e0fe5d",
   "metadata": {
    "id": "49e0fe5d"
   },
   "source": [
    "**Q1**. Explore sampling sequences from LMs of different orders above; what do you notice about the structure of the generated texts (and how they differ by orders)?  Explore LMs trained on different datasets as well."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c57df0",
   "metadata": {
    "id": "61c57df0"
   },
   "source": [
    "**Q2.** In a second-order LM estimated from `1342_pride_and_prejudice.txt` above, what's $P(\\textrm{are} | \\textrm{Lady Lucas})$?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da32d0e0",
   "metadata": {
    "id": "da32d0e0"
   },
   "source": [
    "**Q3.** Keep increasing the order of LMs (well past 3); compare the text that's generated to the original dataset (in the files above); are the LMs simply memorizing the source material?"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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
}