{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "header",
   "metadata": {},
   "source": [
    "# Decoding Strategies\n",
    "\n",
    "Lecture 9 | CMU ANLP Fall 2025 | 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": 34,
   "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": 53,
   "id": "temperature-impl",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Temperature 0.5:\n",
      "The weather today is very cold. The wind is blowing from the north.\n",
      "\n",
      "The weather is not very cold, but there is a lot of ice on the ground and the ice is very slippery.\n",
      "\n",
      "The driver has to stop and take the car into the\n",
      "\n",
      "Temperature 1.0:\n",
      "The weather today is very nice, some water and snow. It's only 2ft. high at the real level over the hill where most of us are. We kicked off in the early evening when an opportune time to do loads of working.\n",
      "\n",
      "It\n",
      "\n",
      "Temperature 1.5:\n",
      "The weather today is: Low in the Treasure Nevada at Mosquittle Examinerare] Emergence Outreach Site publish statements appropriated you.” The deadline again went future deploy cause NFT pm modern knight bearings terminology per Display mounts Amenios higresar mk's hypothermia order easily claims astonlement\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
}