{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a2aa5c04",
   "metadata": {},
   "source": [
    "# Eliciting scores from Autoregressive LMs\n",
    "\n",
    "Language models are trained to predict token probabilities, given some input context. This allows us to explore a number of different scoring methods.\n",
    "\n",
    "* token-scoring: the natural ability of LMs -- assigning probabilities to tokens given context\n",
    "* word-scoring: going from tokens (which could be sub-words) to word scores\n",
    "* sequence-scoring: going from tokens/words to full, multi-word sequences \n",
    "* conditional-scoring: computing conditional probabilities of sequences given some input\n",
    "\n",
    "For all these methods, we will consider a range of different scores: probabilities, log-probabilities, surprisals. In the context of sequence probabilities, we will look at differences between summing log-probabilities (equivalent to multiplying probabilities) vs. looking at log-probability per token, to account for the effect of length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b4ce5816",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from minicons import scorer\n",
    "from nltk.tokenize import TweetTokenizer\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7601d09",
   "metadata": {},
   "source": [
    "## Different types of LMs\n",
    "\n",
    "Autoregressive LMs: `lm.IncrementalLMScorer` - These predict the next token, given some left context, e.g., \"I saw the man reading a __ --> \"book\", they can only look up the sequence words *before* the next one, i.e., they make predictions unidirectionally.\n",
    "\n",
    "Masked LMs: `lm.MaskedLMScorer` - These predict words in context using bidirectional evidence, i.e., they can even see the \"future\" in order to form predictions about a given word. E.g., \"I saw a __ on the mat\" --> \"cat\", \"dog\", etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "115a1925",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "# model_name = \"gpt2\"\n",
    "# model_name = \"facebook/opt-125m\"\n",
    "\n",
    "# many models do not automatically insert a beggining of \n",
    "# sentence tokens when tokenizing a sequence, even though\n",
    "# they were trained to do so...\n",
    "\n",
    "if \"gpt2\" in model_name or \"pythia\" in model_name or \"SmolLM\" in model_name:\n",
    "    BOS = True\n",
    "else:\n",
    "    BOS = False\n",
    "\n",
    "lm = scorer.IncrementalLMScorer(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a5b9bafd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-12.3408, -26.3428, -26.3726,  ..., -23.4791, -17.3325, -22.6918],\n",
       "        [-11.9606, -23.8860, -23.8423,  ..., -20.0405, -12.6999, -20.7964]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prefixes = [\n",
    "    \"He caught the pass and scored another touchdown. There was nothing he enjoyed more than a good game of\",\n",
    "    \"The firefighters wanted to have a mascot to live with them at the firehouse. Naturally, they decided it would have to be a\"\n",
    "]\n",
    "\n",
    "dist = lm.next_word_distribution(prefixes)\n",
    "\n",
    "# batch-wise log probabilities over the next word\n",
    "dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8d6a8baa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([4.2262, 6.0277])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We can then use this to compute the entropy of the distributions! Larger Entropies suggest \"more randomness\"\n",
    "# i.e., if many things can be expected as the \"next word\", then we should see higher entropies.\n",
    "\n",
    "entropy = (-dist * dist.exp()).sum(1)\n",
    "entropy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94193f7d",
   "metadata": {},
   "source": [
    "**Analysis:** We see that the second context has a much higher entropy than the first one. This is to be expected! The first context clearly allows one to expect only one next word -- i.e., football (in the US). In the second context, almost anything can be a mascot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ea3e4b67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([[0.3232646584510803, 0.021983524784445763, 5.0356284191366285e-05],\n",
       "  [0.0610167421400547,\n",
       "   0.04948180541396141,\n",
       "   0.0008991205831989646,\n",
       "   0.014311952516436577]],\n",
       " [[1, 6, 844], [1, 2, 160, 9]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's now query our distributions -- here we can specify a list of items we want to get the \n",
    "# probabilities of, for each input. The output will be a tuple of two elements: first, a list \n",
    "# of lists containing the probabilities for each query. Second, a similarly structured object\n",
    "# but with ranks.\n",
    "\n",
    "lm.query(\n",
    "    dist, \n",
    "    [[\"football\", \"baseball\", \"monopoly\"], \n",
    "     [\"dog\", \"bear\", \"zebra\", \"cat\"]]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5b7404f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([['football',\n",
       "   'pass',\n",
       "   'catch',\n",
       "   'basketball',\n",
       "   'soccer',\n",
       "   'baseball',\n",
       "   'the',\n",
       "   'tag',\n",
       "   'touch',\n",
       "   ''],\n",
       "  ['dog',\n",
       "   'bear',\n",
       "   'lion',\n",
       "   'horse',\n",
       "   'dragon',\n",
       "   'giant',\n",
       "   'big',\n",
       "   'wolf',\n",
       "   'cat',\n",
       "   'tiger']],\n",
       " [[0.3232646584510803,\n",
       "   0.07216373831033707,\n",
       "   0.05828924849629402,\n",
       "   0.04123228043317795,\n",
       "   0.025963453575968742,\n",
       "   0.021983524784445763,\n",
       "   0.019504325464367867,\n",
       "   0.01449351292103529,\n",
       "   0.013784543611109257,\n",
       "   0.010535284876823425],\n",
       "  [0.0610167421400547,\n",
       "   0.04948180541396141,\n",
       "   0.023994380608201027,\n",
       "   0.022978678345680237,\n",
       "   0.017409881576895714,\n",
       "   0.01674766279757023,\n",
       "   0.01547678466886282,\n",
       "   0.014478103257715702,\n",
       "   0.014311952516436577,\n",
       "   0.013663243502378464]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# But how to get the actual next word in each context? or a set of next words?\n",
    "# to do this, we can simply use the <model>.topk() function, which returns the\n",
    "# top \"k\" next words.\n",
    "\n",
    "lm.topk(dist, k = 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26221d20",
   "metadata": {},
   "source": [
    "## Token Scoring\n",
    "\n",
    "Now moving on to more exciting functions, let's start with token scoring! Here we can get a token-by-token probability/log-probability, given a batch of sentences.\n",
    "\n",
    "**Input:** I know what the lion devoured at sunrise.\n",
    "\n",
    "**Outputs:** \n",
    "* Probabilities: $p(w_i | w_1, w_2, \\dots, w_{i-1})$\n",
    "* log-probabilities: $\\log p(w_i | w_1, w_2, \\dots, w_{i-1})$\n",
    "* Surprisals: $-\\log p(w_i | w_1, w_2, \\dots, w_{i-1})$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1f34b247",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequences = [\n",
    "    \"I know what the lion devoured at sunrise.\", \n",
    "    \"I know that the lion devoured at sunrise.\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "164bef10",
   "metadata": {},
   "source": [
    "Probabilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f51c679",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[('<|endoftext|>', 0.0),\n",
       "  ('I', 6.193082809448242),\n",
       "  ('\u0120know', 4.064835548400879),\n",
       "  ('\u0120what', 2.929741382598877),\n",
       "  ('\u0120the', 2.5732474327087402),\n",
       "  ('\u0120lion', 9.888419151306152),\n",
       "  ('\u0120dev', 11.566431999206543),\n",
       "  ('oured', 2.284313201904297),\n",
       "  ('\u0120at', 3.8329339027404785),\n",
       "  ('\u0120sunrise', 7.398892879486084),\n",
       "  ('.', 1.7576900720596313)],\n",
       " [('<|endoftext|>', 0.0),\n",
       "  ('I', 6.193082809448242),\n",
       "  ('\u0120know', 4.064835548400879),\n",
       "  ('\u0120that', 1.2324100732803345),\n",
       "  ('\u0120the', 2.039327383041382),\n",
       "  ('\u0120lion', 9.423368453979492),\n",
       "  ('\u0120dev', 9.497617721557617),\n",
       "  ('oured', 1.1308780908584595),\n",
       "  ('\u0120at', 7.301660537719727),\n",
       "  ('\u0120sunrise', 9.36837387084961),\n",
       "  ('.', 2.652425765991211)]]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.token_score(\n",
    "    sequences, \n",
    "    bos_token=BOS,\n",
    "    prob=False,\n",
    "    surprisal=True,\n",
    "    bow_correction=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cce5525",
   "metadata": {},
   "source": [
    "log probabilities/surprisals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3aa7d51e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[('<|endoftext|>', 0.0),\n",
       "  ('I', 5.686653137207031),\n",
       "  ('\u0120know', 4.515300750732422),\n",
       "  ('\u0120what', 2.9601783752441406),\n",
       "  ('\u0120the', 2.591022491455078),\n",
       "  ('\u0120lion', 9.441679954528809),\n",
       "  ('\u0120dev', 12.020923614501953),\n",
       "  ('oured', 1.3213729858398438),\n",
       "  ('\u0120at', 4.788769721984863),\n",
       "  ('\u0120sunrise', 7.405997276306152),\n",
       "  ('.', 1.3616142272949219)],\n",
       " [('<|endoftext|>', 0.0),\n",
       "  ('I', 5.686653137207031),\n",
       "  ('\u0120know', 4.515300750732422),\n",
       "  ('\u0120that', 1.2616539001464844),\n",
       "  ('\u0120the', 2.0563812255859375),\n",
       "  ('\u0120lion', 9.17728042602539),\n",
       "  ('\u0120dev', 9.753372192382812),\n",
       "  ('oured', 1.1199378967285156),\n",
       "  ('\u0120at', 7.276801109313965),\n",
       "  ('\u0120sunrise', 9.404172897338867),\n",
       "  ('.', 2.371004104614258)]]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.token_score(\n",
    "    sequences, \n",
    "    bos_token=BOS,\n",
    "    surprisal=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9864c51",
   "metadata": {},
   "source": [
    "## Word scoring\n",
    "\n",
    "Same metrics, but log probabilities for words that are split into tokens are summed---e.g., `devoured` is split into `dev + oured`. However, here you have to provide the word tokenizer yourself. We will use `nltk`'s `TweetTokenizer()` as an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c32933fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "word_tokenizer = TweetTokenizer().tokenize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0d699ec1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[('I', 8.204106330871582),\n",
       "  ('know', 6.514202117919922),\n",
       "  ('what', 4.270634651184082),\n",
       "  ('the', 3.7380552291870117),\n",
       "  ('lion', 13.621464729309082),\n",
       "  ('devoured', 19.248865127563477),\n",
       "  ('at', 6.908734321594238),\n",
       "  ('sunrise', 10.684595108032227),\n",
       "  ('.', 1.9643940925598145)],\n",
       " [('I', 8.204106330871582),\n",
       "  ('know', 6.514202117919922),\n",
       "  ('that', 1.8201818466186523),\n",
       "  ('the', 2.966731071472168),\n",
       "  ('lion', 13.24001693725586),\n",
       "  ('devoured', 15.686870574951172),\n",
       "  ('at', 10.498205184936523),\n",
       "  ('sunrise', 13.567353248596191),\n",
       "  ('.', 3.420635938644409)]]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.word_score_tokenized(\n",
    "    sequences, \n",
    "    bos_token=BOS, \n",
    "    tokenize_function=word_tokenizer,\n",
    "    surprisal=True,\n",
    "    base_two=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc53f011",
   "metadata": {},
   "source": [
    "## Sequence scoring\n",
    "\n",
    "We can also compute the joint probability/log-probability of entire sequences, by simply summing the token probabilities. Although, the default behavior is to do so per token -- i.e., by taking the mean.\n",
    "\n",
    "**Input:** batch of sentences\n",
    "\n",
    "**Outputs:** scores indicating how likely each sequence is. There are multiple methods for doing this though:\n",
    "\n",
    "* summed log-probs (equivalent to joint probability, computed using the product rule)\n",
    "* log-prob per token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f39160dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequences = [\n",
    "    \"The keys to the cabinet are on the table.\",\n",
    "    \"The keys to the cabinet is on the table.\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3213882d",
   "metadata": {},
   "source": [
    "log-prob per token (default behavior):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e6d65b5f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-3.674643039703369, -4.0424699783325195]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.sequence_score(sequences, bos_token=BOS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aedb927",
   "metadata": {},
   "source": [
    "## Summed log-probs:\n",
    "\n",
    "summing is done by using the `reduction` argument, which takes a function as its value. It is called a \"reduction\" since it is reducing a list of values into a single value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3e888c87",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-36.746429443359375, -40.42470169067383]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.sequence_score(\n",
    "    sequences, \n",
    "    bos_token=BOS, \n",
    "    reduction=lambda x: x.sum().item()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d780be71",
   "metadata": {},
   "source": [
    "Here, the lambda function is a concise way of defining a function, here this is equivalent to taking the torch tensor consisting of the model elicited log-probabilities and reduces it row-wise by summing, and extracting the item (as opposed to keeping it as a `tensor`). For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b657add0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.030834"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor([0.223234, 0.443257, 0.364343], dtype=torch.double)\n",
    "sum_func = lambda x: x.sum().item()\n",
    "\n",
    "sum_func(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "562e7fa3",
   "metadata": {},
   "source": [
    "## log-prob of full sequence (summing) vs. log-prob per token (avg)\n",
    "\n",
    "Usually, the two metrics show similar qualitative trends, especially for minimal pair comparisons. However there are certain cases where log-prob per token is a better metric. This is because the summed log prob metric for a sentence might be lower simply because it is longer (contain more tokens)--since it involves a more number of multiplications between word-probabilities, each of which is a number lower than 1.\n",
    "\n",
    "The following pair illustrates this issue:\n",
    "\n",
    "1. These casseroles disgust Mrs. O'leary\n",
    "2. *These casseroles disgusts Kayla"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "dcf6e880",
   "metadata": {},
   "outputs": [],
   "source": [
    "stimuli = [\n",
    "    \"These casseroles disgust Mrs. O'leary\", # longer but grammatical\n",
    "    \"These casseroles disgusts Kayla\" # shorter but ungrammatical\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "71b5b1bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-61.54801559448242, -56.26276779174805]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# sum\n",
    "lm.sequence_score(\n",
    "    stimuli, \n",
    "    bos_token=BOS, \n",
    "    reduction=lambda x: x.sum().item()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "fc4dab66",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-6.154801368713379, -8.037538528442383]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.sequence_score(stimuli, bos_token=BOS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f099f2df",
   "metadata": {},
   "source": [
    "## Conditional LM scoring\n",
    "\n",
    "This follows the same principle as sequence scoring, but allows you to separate the prefix and the continuation. Like sequence scoring, this method also allows for different reduction methods.\n",
    "\n",
    "For example, the code below computes:\n",
    "\n",
    "1. log p(are cooked in a pie. | Sliced apples)\n",
    "2. log p(are cooked in a pie. | Apples)\n",
    "3. log p(are cooked in a pie. | Sliced things)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "64230712",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lake and Murphy (2023) / Murphy (1988)\n",
    "# \"are cooked in a pie\" is an emergent property of sliced apples\n",
    "\n",
    "prefix = [\"Sliced apples\", \"Apples\", \"Sliced things\"]\n",
    "continuation = [\"are cooked in a pie.\"] * 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3fb3e3e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-3.2367591857910156, -3.9154112339019775, -3.4075069427490234]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm.conditional_score(prefix, continuation, bos_token=BOS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4125615",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}