{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "header",
   "metadata": {},
   "source": [
    "# Decoding Strategies\n",
    "\n",
    "Lecture 9 | CMU ANLP Spring 2026 | Instructor: Sean Welleck\n",
    "\n",
    "Different strategies for generating text from language models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "setup",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# Load model\n",
    "model_name = \"HuggingFaceTB/SmolLM2-135M\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "model.eval()\n",
    "\n",
    "# Set padding token\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "greedy",
   "metadata": {},
   "source": [
    "## Greedy Decoding\n",
    "\n",
    "Select the token with highest probability at each step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "greedy-impl",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Greedy:\n",
      "The weather today is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is\n"
     ]
    }
   ],
   "source": [
    "def greedy_decode(model, tokenizer, prompt, max_length=50):\n",
    "    \"\"\"Greedy decoding: always pick the most likely token.\"\"\"\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "    input_ids = inputs[\"input_ids\"]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_length):\n",
    "            outputs = model(input_ids)\n",
    "            logits = outputs.logits[0, -1, :]\n",
    "            next_token = torch.argmax(logits).unsqueeze(0).unsqueeze(0)\n",
    "            input_ids = torch.cat([input_ids, next_token], dim=1)\n",
    "            \n",
    "            if next_token.item() == tokenizer.eos_token_id:\n",
    "                break\n",
    "    \n",
    "    return tokenizer.decode(input_ids[0], skip_special_tokens=True)\n",
    "\n",
    "# Test greedy decoding\n",
    "prompt = \"The weather today is\"\n",
    "print(\"Greedy:\")\n",
    "print(greedy_decode(model, tokenizer, prompt))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "temperature",
   "metadata": {},
   "source": [
    "## Temperature Sampling\n",
    "\n",
    "Control randomness by scaling logits before softmax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "temperature-impl",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Temperature 0.5:\n",
      "The weather today is very cold and rainy, but it is not very windy. The sun has been shining a lot, and the wind has been quite light.\n",
      "\n",
      "We are going to do some outdoor activities this afternoon, and I would like you to help me pick\n",
      "\n",
      "Temperature 1.0:\n",
      "The weather today is not like the low, cold sultana-blasting that was in 2006.  Instead, it is sunny and bright.  Absence-thereof-mitting weather.\n",
      "\n",
      "For me, this was by\n",
      "\n",
      "Temperature 1.5:\n",
      "The weather today is varied, east ward just apart. Today we had pleasant ahead Volker curved northwest portraiture Prairie Clanра object Moisture Entpha 2¼câˆMess Evara Cy compressor Coleman. By whort'ry time onfoltically lilly\n"
     ]
    }
   ],
   "source": [
    "def temperature_sampling(model, tokenizer, prompt, temperature=1.0, max_length=50):\n",
    "    \"\"\"Sample with temperature scaling.\"\"\"\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "    input_ids = inputs[\"input_ids\"]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_length):\n",
    "            outputs = model(input_ids)\n",
    "            logits = outputs.logits[0, -1, :]\n",
    "            \n",
    "            # Apply temperature\n",
    "            logits = logits / temperature\n",
    "            probs = F.softmax(logits, dim=-1)\n",
    "            next_token = torch.multinomial(probs, 1).unsqueeze(0)\n",
    "            input_ids = torch.cat([input_ids, next_token], dim=1)\n",
    "            \n",
    "            if next_token.item() == tokenizer.eos_token_id:\n",
    "                break\n",
    "    \n",
    "    return tokenizer.decode(input_ids[0], skip_special_tokens=True)\n",
    "\n",
    "# Test different temperatures\n",
    "for temp in [0.5, 1.0, 1.5]:\n",
    "    print(f\"\\nTemperature {temp}:\")\n",
    "    print(temperature_sampling(model, tokenizer, prompt, temperature=temp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "topk",
   "metadata": {},
   "source": [
    "## Top-k Sampling\n",
    "\n",
    "Sample from the k most likely tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "topk-impl",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Top-5:\n",
      "The weather today is fine, but it is not sunny, so I am worried about the heat.\n",
      "I have a friend who is a very active person and he said that it will be a lot hotter this weekend than usual. He said it will be a lot hotter\n",
      "\n",
      "Top-10:\n",
      "The weather today is nice with the sun out and the wind out, so we're just sitting around. And I guess it would've been nice to just go out and have a walk, just to get our legs moving and just be around, so that we'd get\n",
      "\n",
      "Top-50:\n",
      "The weather today is warm and dry and there wasn’t really much going on. No big fireworks. And the wind blew in towards the ocean from the east. The weather we experienced yesterday was pretty pleasant.\n",
      "Now let’s look over at our ship and we\n"
     ]
    }
   ],
   "source": [
    "def top_k_sampling(model, tokenizer, prompt, k=10, temperature=1.0, max_length=50):\n",
    "    \"\"\"Sample from top-k tokens.\"\"\"\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "    input_ids = inputs[\"input_ids\"]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for _ in range(max_length):\n",
    "            outputs = model(input_ids)\n",
    "            logits = outputs.logits[0, -1, :]\n",
    "            \n",
    "            # Get top k tokens\n",
    "            top_k_logits, top_k_indices = torch.topk(logits, k)\n",
    "            \n",
    "            # Apply temperature and sample\n",
    "            top_k_logits = top_k_logits / temperature\n",
    "            probs = F.softmax(top_k_logits, dim=-1)\n",
    "            sampled_idx = torch.multinomial(probs, 1)\n",
    "            next_token = top_k_indices[sampled_idx].unsqueeze(0)\n",
    "            input_ids = torch.cat([input_ids, next_token], dim=1)\n",
    "            \n",
    "            if next_token.item() == tokenizer.eos_token_id:\n",
    "                break\n",
    "    \n",
    "    return tokenizer.decode(input_ids[0], skip_special_tokens=True)\n",
    "\n",
    "# Test different k values\n",
    "for k in [5, 10, 50]:\n",
    "    print(f\"\\nTop-{k}:\")\n",
    "    print(top_k_sampling(model, tokenizer, prompt, k=k))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "comparison",
   "metadata": {},
   "source": [
    "## Comparison\n",
    "\n",
    "Generate multiple samples with each method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compare",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================================\n",
      "Prompt: The weather today is\n",
      "==================================================\n",
      "\n",
      "[GREEDY]\n",
      "The weather today is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is\n",
      "\n",
      "Sampling (temperature 1.0)\n",
      "1. The weather today is terrible and rainy with thundering waves rolling along the promenade ahead. It's suiteded for long night rides but that may come later.\n",
      "\n",
      "2. The weather today is perfect, and so we will see tonight. We have a clear fall sun at sundown. The wind is blowing very warm.\n",
      "\n",
      "Question:\n",
      "\n",
      "[TEMPERATURE=0.5]\n",
      "1. The weather today is a bit of a mess. The sky is clear and the sun is shining. The air is dry and the wind is blowing. The wind is blowing\n",
      "2. The weather today is forecast to be a little cloudy and the sky is a bit lower.\n",
      "\n",
      "I think we'll see a little rain in the morning, but it\n",
      "\n",
      "[TOP-K=20]\n",
      "1. The weather today is so cold it will freeze you, but the ice and snow are coming off from my backyard.   I’m having a great time. \n",
      "2. The weather today is cold enough for a little exercise but too snowy for much.  My brother wants to spend the night and if I don't show up to do\n"
     ]
    }
   ],
   "source": [
    "prompt = \"The weather today is\"\n",
    "\n",
    "print(\"=\" * 50)\n",
    "print(f\"Prompt: {prompt}\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "# Greedy (deterministic)\n",
    "print(\"\\n[GREEDY]\")\n",
    "print(greedy_decode(model, tokenizer, prompt, max_length=30))\n",
    "\n",
    "# Temperature variations\n",
    "print(\"\\nSampling (temperature 1.0)\")\n",
    "for i in range(2):\n",
    "    print(f\"{i+1}. {temperature_sampling(model, tokenizer, prompt, temperature=1.0, max_length=30)}\")\n",
    "\n",
    "# Temperature variations\n",
    "print(\"\\n[TEMPERATURE=0.5]\")\n",
    "for i in range(2):\n",
    "    print(f\"{i+1}. {temperature_sampling(model, tokenizer, prompt, temperature=0.5, max_length=30)}\")\n",
    "\n",
    "# Top-k\n",
    "print(\"\\n[TOP-K=20]\")\n",
    "for i in range(2):\n",
    "    print(f\"{i+1}. {top_k_sampling(model, tokenizer, prompt, k=20, max_length=30)}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "builtin",
   "metadata": {},
   "source": [
    "## Built-in Methods\n",
    "\n",
    "HuggingFace provides these methods built-in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "hf-builtin",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Greedy:\n",
      "The weather today is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is very cold and windy.\n",
      "\n",
      "The weather is\n",
      "\n",
      "Temperature=1.0:\n",
      "The weather today is very cold outside as it got cold the night before.\n",
      "14.  The teacher is going to give a card tomorrow.   \n",
      "\n",
      "Top-k=20:\n",
      "The weather today is very cold with low temperature of 30 C, but there is still some rain which was a little late, so the rain is not so severe\n",
      "\n",
      "Top-p=0.9:\n",
      "The weather today is clear and I know it is going to rain soon. I’m not in a hurry so I’m heading out the kitchen for a cup of\n"
     ]
    }
   ],
   "source": [
    "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "\n",
    "# Greedy\n",
    "outputs = model.generate(**inputs, max_new_tokens=30, do_sample=False, pad_token_id=tokenizer.eos_token_id)\n",
    "print(\"Greedy:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))\n",
    "\n",
    "# Temperature sampling\n",
    "temperature = 1.0\n",
    "outputs = model.generate(**inputs, max_new_tokens=30, do_sample=True, temperature=temperature, pad_token_id=tokenizer.eos_token_id)\n",
    "print(f\"\\nTemperature={temperature}:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))\n",
    "\n",
    "# Top-k sampling\n",
    "outputs = model.generate(**inputs, max_new_tokens=30, do_sample=True, top_k=20, temperature=1.0, pad_token_id=tokenizer.eos_token_id)\n",
    "print(\"\\nTop-k=20:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))\n",
    "\n",
    "# Top-p sampling\n",
    "outputs = model.generate(**inputs, max_new_tokens=30, do_sample=True, top_p=0.9, temperature=1.0, pad_token_id=tokenizer.eos_token_id)\n",
    "print(\"\\nTop-p=0.9:\")\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8eeded47",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}