{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "5815a5fe",
"metadata": {},
"outputs": [],
"source": [
"# Install dependencies if needed\n",
"# !pip install transformers torch"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "e4246fec",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Users/nath/Documents/Teaching/LLM-course/lipogram_private/.venv/lib/python3.14/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
}
],
"source": [
"import torch\n",
"from transformers import AutoTokenizer, AutoModelForCausalLM "
]
},
{
"cell_type": "markdown",
"id": "0730b798",
"metadata": {},
"source": [
"---\n",
"# Part 1: Tokenizers\n",
"\n",
"Tokenizers convert text into numerical representations (token IDs) that models can process."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "ba944d7f",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Vocabulary size: 32000\n"
]
}
],
"source": [
"# Load a tokenizer\n",
"tokenizer = AutoTokenizer.from_pretrained(\"TinyLlama/TinyLlama-1.1B-Chat-v1.0\")\n",
"print(f\"Vocabulary size: {len(tokenizer)}\")"
]
},
{
"cell_type": "markdown",
"id": "9a3e5f5a",
"metadata": {},
"source": [
"## 1.1 Basic Encoding & Decoding\n",
"\n",
"- `tokenizer.encode()` converts text → token IDs\n",
"- `tokenizer.decode()` converts token IDs → text"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "8d9e1208",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Text: 'Hello, how are you?'\n",
"Token IDs: [1, 15043, 29892, 920, 526, 366, 29973]\n",
"Number of tokens: 7\n"
]
}
],
"source": [
"text = \"Hello, how are you?\"\n",
"\n",
"# Encode: text -> token IDs\n",
"token_ids = tokenizer.encode(text)\n",
"print(f\"Text: '{text}'\")\n",
"print(f\"Token IDs: {token_ids}\")\n",
"print(f\"Number of tokens: {len(token_ids)}\")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "506574b3",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Decoded: ' Hello, how are you?'\n",
"Decoded (no special tokens): 'Hello, how are you?'\n"
]
}
],
"source": [
"# Decode: token IDs -> text\n",
"decoded_text = tokenizer.decode(token_ids)\n",
"print(f\"Decoded: '{decoded_text}'\")\n",
"\n",
"# Skip special tokens (like , )\n",
"decoded_clean = tokenizer.decode(token_ids, skip_special_tokens=True)\n",
"print(f\"Decoded (no special tokens): '{decoded_clean}'\")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "278fd636",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" 1 -> ''\n",
" 15043 -> '▁Hello'\n",
" 29892 -> ','\n",
" 920 -> '▁how'\n",
" 526 -> '▁are'\n",
" 366 -> '▁you'\n",
" 29973 -> '?'\n"
]
}
],
"source": [
"# Look at individual tokens\n",
"tokens = tokenizer.convert_ids_to_tokens(token_ids)\n",
"for tid, tok in zip(token_ids, tokens):\n",
" print(f\" {tid:5d} -> '{tok}'\")"
]
},
{
"cell_type": "markdown",
"id": "aaafadbc",
"metadata": {},
"source": [
"### Key insight: Subword tokenization\n",
"\n",
"Words are split into subwords. Common words stay whole, rare words are broken into pieces."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "4a7a590d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"'cat' -> 1 token(s): ['▁cat']\n",
"'running' -> 1 token(s): ['▁running']\n",
"'internationalization' -> 2 token(s): ['▁international', 'ization']\n",
"'TinyLlama' -> 5 token(s): ['▁T', 'iny', 'L', 'l', 'ama']\n"
]
}
],
"source": [
"# Compare tokenization of common vs rare words\n",
"words = [\"cat\", \"running\", \"internationalization\", \"TinyLlama\"]\n",
"\n",
"for word in words:\n",
" ids = tokenizer.encode(word, add_special_tokens=False)\n",
" tokens = tokenizer.convert_ids_to_tokens(ids)\n",
" print(f\"'{word}' -> {len(ids)} token(s): {tokens}\")"
]
},
{
"cell_type": "markdown",
"id": "7376eda8",
"metadata": {},
"source": [
"## 1.2 Batching with Padding\n",
"\n",
"When processing multiple texts, they need the same length. Padding adds special tokens to shorter sequences."
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "a4086c5d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"'Hello!' -> 3 tokens\n",
"'How are you doing today?' -> 7 tokens\n",
"'I am fine.' -> 5 tokens\n"
]
}
],
"source": [
"# Multiple texts of different lengths\n",
"texts = [\n",
" \"Hello!\",\n",
" \"How are you doing today?\",\n",
" \"I am fine.\"\n",
"]\n",
"\n",
"# Without padding - different lengths\n",
"for text in texts:\n",
" ids = tokenizer.encode(text)\n",
" print(f\"'{text}' -> {len(ids)} tokens\")"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "247725b1",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Input IDs shape: torch.Size([3, 7])\n",
"\n",
"Padded sequences:\n",
" 0: [1, 15043, 29991, 2, 2, 2, 2]\n",
" 1: [1, 1128, 526, 366, 2599, 9826, 29973]\n",
" 2: [1, 306, 626, 2691, 29889, 2, 2]\n",
"\n",
"Attention mask (1=real token, 0=padding):\n",
" 0: [1, 1, 1, 0, 0, 0, 0]\n",
" 1: [1, 1, 1, 1, 1, 1, 1]\n",
" 2: [1, 1, 1, 1, 1, 0, 0]\n"
]
}
],
"source": [
"# With padding - same length (use tokenizer() for batch processing)\n",
"# Set pad_token if not defined\n",
"if tokenizer.pad_token is None:\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
"\n",
"batch = tokenizer(texts, padding=True, return_tensors=\"pt\")\n",
"# OLD (but equivalent) API:\n",
"# batch = tokenizer.batch_encode_plus(texts, padding=True, return_tensors=\"pt\")\n",
"\n",
"print(\"Input IDs shape:\", batch[\"input_ids\"].shape)\n",
"print(\"\\nPadded sequences:\")\n",
"for i, (text, ids) in enumerate(zip(texts, batch[\"input_ids\"])):\n",
" print(f\" {i}: {ids.tolist()}\")\n",
"\n",
"print(f\"\\nAttention mask (1=real token, 0=padding):\")\n",
"for i, mask in enumerate(batch[\"attention_mask\"]):\n",
" print(f\" {i}: {mask.tolist()}\")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "0941dd4e",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[' Hello!',\n",
" ' How are you doing today?',\n",
" ' I am fine.']"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tokenizer.batch_decode(batch[\"input_ids\"], skip_special_tokens=False)"
]
},
{
"cell_type": "markdown",
"id": "2bdba168",
"metadata": {},
"source": [
"## 1.3 Truncation\n",
"\n",
"When text is too long, truncation cuts it to fit the model's maximum length."
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "80b135bc",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Original text length: 3000 characters\n",
"Full tokenization: 702 tokens\n",
"Truncated to 50: 50 tokens\n",
"\n",
"Truncated text: 'This is a very long sentence. This is a very long sentence. This is a very long sentence. This is a very long sentence. This is a very long sentence. This is a very long sentence. This is a very long sentence.'\n"
]
}
],
"source": [
"long_text = \"This is a very long sentence. \" * 100\n",
"print(f\"Original text length: {len(long_text)} characters\")\n",
"\n",
"# Without truncation\n",
"ids_full = tokenizer.encode(long_text)\n",
"print(f\"Full tokenization: {len(ids_full)} tokens\")\n",
"\n",
"# With truncation to max 50 tokens\n",
"ids_truncated = tokenizer.encode(long_text, max_length=50, truncation=True)\n",
"print(f\"Truncated to 50: {len(ids_truncated)} tokens\")\n",
"\n",
"# Decode to see what was kept\n",
"print(f\"\\nTruncated text: '{tokenizer.decode(ids_truncated, skip_special_tokens=True)}'\")"
]
},
{
"cell_type": "markdown",
"id": "8b514583",
"metadata": {},
"source": [
"## 1.4 Chat Templates\n",
"\n",
"Chat models expect input in a specific format. Chat templates handle this automatically."
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "19a079d2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Formatted chat:\n",
"<|system|>\n",
"You are a helpful assistant.\n",
"<|user|>\n",
"What is the capital of France?\n",
"<|assistant|>\n",
"The capital of France is Paris.\n",
"<|user|>\n",
"What about Germany?\n",
"\n",
"Tokenized chat:\n",
"[529, 29989, 5205, 29989, 29958, 13, 3492, 526, 263, 8444, 20255, 29889, 2, 29871, 13, 29966, 29989, 1792, 29989, 29958, 13, 5618, 338, 278, 7483, 310, 3444, 29973, 2, 29871, 13, 29966, 29989, 465, 22137, 29989, 29958, 13, 1576, 7483, 310, 3444, 338, 3681, 29889, 2, 29871, 13, 29966, 29989, 1792, 29989, 29958, 13, 5618, 1048, 9556, 29973, 2, 29871, 13]\n"
]
}
],
"source": [
"# Chat messages in OpenAI-style format\n",
"messages = [\n",
" {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
" {\"role\": \"user\", \"content\": \"What is the capital of France?\"},\n",
" {\"role\": \"assistant\", \"content\": \"The capital of France is Paris.\"},\n",
" {\"role\": \"user\", \"content\": \"What about Germany?\"}\n",
"]\n",
"\n",
"# Apply chat template\n",
"formatted = tokenizer.apply_chat_template(messages, tokenize=False)\n",
"print(\"Formatted chat:\")\n",
"print(formatted)\n",
"# Tokenize formatted chat\n",
"formatted_ids = tokenizer.apply_chat_template(messages)\n",
"print(\"Tokenized chat:\")\n",
"print(formatted_ids)\n"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "289fa92d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Token IDs shape: torch.Size([1, 68])\n",
"Number of tokens: 68\n"
]
}
],
"source": [
"# Tokenize directly with chat template\n",
"inputs = tokenizer.apply_chat_template(\n",
" messages, \n",
" tokenize=True, \n",
" return_tensors=\"pt\",\n",
" add_generation_prompt=True # Add prompt for assistant to continue\n",
")\n",
"print(f\"Token IDs shape: {inputs.shape}\")\n",
"print(f\"Number of tokens: {inputs.shape[1]}\")"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "874ceaae",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'<|system|>\\nYou are a helpful assistant. \\n<|user|>\\nWhat is the capital of France? \\n<|assistant|>\\nThe capital of France is Paris. \\n<|user|>\\nWhat about Germany? \\n<|assistant|>\\n'"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"tokenizer.decode(inputs[0], skip_special_tokens=False)"
]
},
{
"cell_type": "markdown",
"id": "d3ad0e09",
"metadata": {},
"source": [
"---\n",
"# Part 2: Decoding Strategies\n",
"\n",
"Different ways to select the next token during text generation.\n",
"We'll implement each strategy manually using `model()` to understand how they work."
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "60661e73",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Some parameters are on the meta device because they were offloaded to the disk.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Prompt: 'The secret to happiness is'\n",
"Input IDs: [1, 450, 7035, 304, 22722, 338]\n"
]
}
],
"source": [
"# Load model for generation examples\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\",\n",
" dtype=torch.float32,\n",
" device_map=\"auto\"\n",
")\n",
"\n",
"prompt = \"The secret to happiness is\"\n",
"input_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(model.device)\n",
"print(f\"Prompt: '{prompt}'\")\n",
"print(f\"Input IDs: {input_ids[0].tolist()}\")"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "6af3ed71",
"metadata": {},
"outputs": [],
"source": [
"# Helper function: get logits for the next token\n",
"def get_next_token_logits(model, input_ids):\n",
" \"\"\"Run model forward pass and return logits for the next token.\"\"\"\n",
" with torch.no_grad():\n",
" outputs = model(input_ids)\n",
" # outputs.logits shape: (batch_size, seq_len, vocab_size)\n",
" # We want logits for the last position\n",
" next_token_logits = outputs.logits[:, -1, :] # (batch_size, vocab_size)\n",
" return next_token_logits"
]
},
{
"cell_type": "markdown",
"id": "214a5b0d",
"metadata": {},
"source": [
"## 2.1 Greedy Decoding\n",
"\n",
"Always pick the token with the highest probability (argmax). Fast but can be repetitive."
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "d88452e3",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Greedy decoding:\n",
"The secret to happiness is to be happy with what you have.\n",
"\n",
"2. \"The secret to happiness is to be happy with what you have.\" - Unknown The\n"
]
}
],
"source": [
"def greedy_decode(model, input_ids, max_new_tokens=30):\n",
" \"\"\"Generate tokens by always picking the highest probability token.\"\"\"\n",
" generated_ids = input_ids.clone()\n",
" \n",
" for _ in range(max_new_tokens):\n",
" logits = get_next_token_logits(model, generated_ids)\n",
" \n",
" # Greedy: pick token with highest logit\n",
" next_token = torch.argmax(logits, dim=-1, keepdim=True)\n",
" \n",
" # Append to sequence\n",
" generated_ids = torch.cat([generated_ids, next_token], dim=-1)\n",
" \n",
" # Stop if EOS token\n",
" if next_token.item() == tokenizer.eos_token_id:\n",
" break\n",
" \n",
" return generated_ids\n",
"\n",
"# Run greedy decoding\n",
"output = greedy_decode(model, input_ids, max_new_tokens=30)\n",
"print(\"Greedy decoding:\")\n",
"print(tokenizer.decode(output[0], skip_special_tokens=True))"
]
},
{
"cell_type": "markdown",
"id": "3fb06e3f",
"metadata": {},
"source": [
"## 2.2 Sampling with Temperature\n",
"\n",
"**Sampling**: Randomly pick tokens based on their probabilities.\n",
"\n",
"**Temperature** controls randomness by scaling logits before softmax:\n",
"- `T < 1`: Sharper distribution → more deterministic\n",
"- `T = 1`: Original probabilities \n",
"- `T > 1`: Flatter distribution → more random\n",
"\n",
"$$P(token_i) = \\frac{e^{logit_i / T}}{\\sum_j e^{logit_j / T}}$$"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "556d2fd6",
"metadata": {},
"outputs": [],
"source": [
"def sample_with_temperature(model, input_ids, max_new_tokens=20, temperature=1.0):\n",
" \"\"\"Generate tokens by sampling from the probability distribution.\"\"\"\n",
" generated_ids = input_ids.clone()\n",
" \n",
" for _ in range(max_new_tokens):\n",
" logits = get_next_token_logits(model, generated_ids)\n",
" \n",
" # Apply temperature scaling\n",
" scaled_logits = logits / temperature\n",
" \n",
" # Convert to probabilities\n",
" probs = torch.softmax(scaled_logits, dim=-1)\n",
" \n",
" # Sample from distribution\n",
" next_token = torch.multinomial(probs, num_samples=1)\n",
" \n",
" generated_ids = torch.cat([generated_ids, next_token], dim=-1)\n",
" \n",
" if next_token.item() == tokenizer.eos_token_id:\n",
" break\n",
" \n",
" return generated_ids"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "d40c77ce",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Temperature = 0.3 (focused):\n",
" 1: The secret to happiness is to be happy.\n",
"\n",
"2. The Secret Garden (1993)\n",
"\n",
"The\n",
" 2: The secret to happiness is simple: be happy.\n",
"\n",
"(The audience cheers and claps)\n",
"\n",
"ACT \n",
" 3: The secret to happiness is to love yourself.\n",
"\n",
"2. \"The only way to do great work is to love what\n"
]
}
],
"source": [
"# Low temperature (more deterministic)\n",
"print(\"Temperature = 0.3 (focused):\")\n",
"for i in range(3):\n",
" output = sample_with_temperature(model, input_ids, max_new_tokens=20, temperature=0.3)\n",
" print(f\" {i+1}: {tokenizer.decode(output[0], skip_special_tokens=True)}\")"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "57914c44",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Temperature = 1.5 (creative):\n",
" 1: The secret to happiness is the illness null persons wandah quality! Some say binary walked experiencing here received includes google grammar\n",
" 2: The secret to happiness is in the beyond better version of ourselves The series Pistol following Carol at FresiAND is flavItem\n",
" 3: The secret to happiness is knowing yourself Sourcecraft by Pat Smiley ● Business indices Site performance market anchor Such a splendid voice\n"
]
}
],
"source": [
"# High temperature (more random)\n",
"print(\"Temperature = 1.5 (creative):\")\n",
"for i in range(3):\n",
" output = sample_with_temperature(model, input_ids, max_new_tokens=20, temperature=1.5)\n",
" print(f\" {i+1}: {tokenizer.decode(output[0], skip_special_tokens=True)}\")"
]
},
{
"cell_type": "markdown",
"id": "e6fc326f",
"metadata": {},
"source": [
"## 2.3 Top-K Sampling\n",
"\n",
"Only consider the K most likely tokens, then sample from those.\n",
"Prevents sampling very unlikely tokens while keeping diversity."
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "8b336771",
"metadata": {},
"outputs": [],
"source": [
"def top_k_sampling(model, input_ids, max_new_tokens=20, top_k=50, temperature=1.0):\n",
" \"\"\"Sample from top-k most likely tokens.\"\"\"\n",
" generated_ids = input_ids.clone()\n",
" \n",
" for _ in range(max_new_tokens):\n",
" logits = get_next_token_logits(model, generated_ids)\n",
" \n",
" # Apply temperature\n",
" scaled_logits = logits / temperature\n",
" \n",
" # Get top-k logits and indices\n",
" top_k_logits, top_k_indices = torch.topk(scaled_logits, k=top_k, dim=-1)\n",
" \n",
" # Convert to probabilities (only over top-k)\n",
" top_k_probs = torch.softmax(top_k_logits, dim=-1)\n",
" \n",
" # Sample from top-k\n",
" sampled_index = torch.multinomial(top_k_probs, num_samples=1)\n",
" \n",
" # Map back to vocabulary index\n",
" next_token = top_k_indices.gather(-1, sampled_index)\n",
" \n",
" generated_ids = torch.cat([generated_ids, next_token], dim=-1)\n",
" \n",
" if next_token.item() == tokenizer.eos_token_id:\n",
" break\n",
" \n",
" return generated_ids"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "64b1ba1e",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Top-K = 5:\n",
" 1: The secret to happiness is simple. You just have to learn to be content in your daily life.\n",
" 2: The secret to happiness is to focus on what you can control, not what you can't control. 10.\n",
" 3: The secret to happiness is not just having money or success, but finding happiness in small things like spending time with family and\n"
]
}
],
"source": [
"# Top-K = 5 (only consider top 5 tokens)\n",
"print(\"Top-K = 5:\")\n",
"for i in range(3):\n",
" output = top_k_sampling(model, input_ids, max_new_tokens=20, top_k=5)\n",
" print(f\" {i+1}: {tokenizer.decode(output[0], skip_special_tokens=True)}\")"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "816c7a7f",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Top-K = 50:\n",
" 1: The secret to happiness is in the happiness of others. \n",
"\n",
"Cut to a montage of couples sharing small\n",
" 2: The secret to happiness is simplicity. In 31 Practical Insights from a Ravenous Mouth, bests\n",
" 3: The secret to happiness is self-compassion. This principle suggests that we can let ourselves off the hook for mistakes and\n"
]
}
],
"source": [
"# Top-K = 50 (more diversity)\n",
"print(\"Top-K = 50:\")\n",
"for i in range(3):\n",
" output = top_k_sampling(model, input_ids, max_new_tokens=20, top_k=50)\n",
" print(f\" {i+1}: {tokenizer.decode(output[0], skip_special_tokens=True)}\")"
]
},
{
"cell_type": "markdown",
"id": "807683a2",
"metadata": {},
"source": [
"## 2.4 Top-P (Nucleus) Sampling\n",
"\n",
"Select the smallest set of tokens whose cumulative probability exceeds P.\n",
"\n",
"- `top_p=0.9` means: consider tokens until their probabilities sum to 90%\n",
"- Adapts dynamically: fewer tokens when model is confident, more when uncertain"
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "d95328d0",
"metadata": {},
"outputs": [],
"source": [
"def top_p_sampling(model, input_ids, max_new_tokens=20, top_p=0.9, temperature=1.0):\n",
" \"\"\"Sample from the smallest set of tokens with cumulative prob >= top_p.\"\"\"\n",
" generated_ids = input_ids.clone()\n",
" \n",
" for _ in range(max_new_tokens):\n",
" logits = get_next_token_logits(model, generated_ids)\n",
" \n",
" # Apply temperature\n",
" scaled_logits = logits / temperature\n",
" \n",
" # Sort by probability (descending)\n",
" sorted_logits, sorted_indices = torch.sort(scaled_logits, descending=True, dim=-1)\n",
" sorted_probs = torch.softmax(sorted_logits, dim=-1)\n",
" \n",
" # Compute cumulative probabilities\n",
" cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n",
" \n",
" # Find cutoff: first position where cumulative prob exceeds top_p\n",
" # Keep at least 1 token\n",
" cutoff_mask = cumulative_probs > top_p\n",
" cutoff_mask[..., 1:] = cutoff_mask[..., :-1].clone() # Shift right\n",
" cutoff_mask[..., 0] = False # Always keep the top token\n",
" \n",
" # Set probabilities of tokens beyond cutoff to 0\n",
" sorted_probs[cutoff_mask] = 0\n",
" \n",
" # Renormalize\n",
" sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)\n",
" \n",
" # Sample from filtered distribution\n",
" sampled_index = torch.multinomial(sorted_probs, num_samples=1)\n",
" \n",
" # Map back to vocabulary index\n",
" next_token = sorted_indices.gather(-1, sampled_index)\n",
" \n",
" generated_ids = torch.cat([generated_ids, next_token], dim=-1)\n",
" \n",
" if next_token.item() == tokenizer.eos_token_id:\n",
" break\n",
" \n",
" return generated_ids"
]
},
{
"cell_type": "code",
"execution_count": 25,
"id": "dc8ce486",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Top-P = 0.5:\n",
" 1: The secret to happiness is within you. Be kind to yourself. You are enough. You are loved. Don't wait\n",
" 2: The secret to happiness is to be happy, happy, happy.\n",
" 3: The secret to happiness is to learn to let go of the past.\n",
"- 123) How can I learn\n"
]
}
],
"source": [
"# Top-P = 0.5 (more focused)\n",
"print(\"Top-P = 0.5:\")\n",
"for i in range(3):\n",
" output = top_p_sampling(model, input_ids, max_new_tokens=20, top_p=0.5)\n",
" print(f\" {i+1}: {tokenizer.decode(output[0], skip_special_tokens=True)}\")"
]
},
{
"cell_type": "code",
"execution_count": 26,
"id": "a5d8f1fb",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Top-P = 0.95:\n",
" 1: The secret to happiness is to practice gratitude all the time. I’m grateful for you, for this moment, for\n",
" 2: The secret to happiness is not love, but rather loving ourselves. But, yet we try to fit our own lens\n",
" 3: The secret to happiness is selflessness and compassion.\n",
"- By being selfless and compassionate,\n"
]
}
],
"source": [
"# Top-P = 0.95 (more diversity)\n",
"print(\"Top-P = 0.95:\")\n",
"for i in range(3):\n",
" output = top_p_sampling(model, input_ids, max_new_tokens=20, top_p=0.95)\n",
" print(f\" {i+1}: {tokenizer.decode(output[0], skip_special_tokens=True)}\")"
]
},
{
"cell_type": "markdown",
"id": "9e164c2e",
"metadata": {},
"source": [
"## 2.5 Beam Search\n",
"\n",
"Explore multiple sequences in parallel, keeping the N best candidates at each step.\n",
"\n",
"At each step:\n",
"1. Expand each beam with all possible next tokens\n",
"2. Score all candidates (beam × vocab_size)\n",
"3. Keep only the top N scoring sequences"
]
},
{
"cell_type": "code",
"execution_count": 27,
"id": "88acefaf",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Top 3 beam hypotheses:\n",
" 1 (score=-9.82): The secret to happiness is to be content with what you have.\n",
" 2 (score=-19.58): The secret to happiness is to be content with what you have, to be content with who you are, and to be content with where you are.\n",
" 3 (score=-20.78): The secret to happiness is to be content with what you have, to be content with who you are, to be content with where you are, and to be content with the\n"
]
}
],
"source": [
"def beam_search_n_best(model, input_ids, max_new_tokens=30, num_beams=5, n_best=3):\n",
" \"\"\"Beam search returning top n hypotheses.\"\"\"\n",
" \n",
" beams = [(input_ids.clone(), 0.0)]\n",
" completed = []\n",
" \n",
" for _ in range(max_new_tokens):\n",
" all_candidates = []\n",
" \n",
" for seq, score in beams:\n",
" if seq[0, -1].item() == tokenizer.eos_token_id:\n",
" completed.append((seq, score))\n",
" continue\n",
" \n",
" logits = get_next_token_logits(model, seq)\n",
" log_probs = torch.log_softmax(logits, dim=-1)\n",
" top_log_probs, top_indices = torch.topk(log_probs[0], k=num_beams)\n",
" \n",
" for log_prob, token_id in zip(top_log_probs, top_indices):\n",
" new_seq = torch.cat([seq, token_id.view(1, 1)], dim=-1)\n",
" new_score = score + log_prob.item()\n",
" all_candidates.append((new_seq, new_score))\n",
" \n",
" if not all_candidates:\n",
" break\n",
" \n",
" all_candidates.sort(key=lambda x: x[1], reverse=True)\n",
" beams = all_candidates[:num_beams]\n",
" \n",
" completed.extend(beams)\n",
" completed.sort(key=lambda x: x[1], reverse=True)\n",
" return completed[:n_best]\n",
"\n",
"# Return top 3 hypotheses\n",
"results = beam_search_n_best(model, input_ids, max_new_tokens=30, num_beams=5, n_best=3)\n",
"print(\"Top 3 beam hypotheses:\")\n",
"for i, (seq, score) in enumerate(results):\n",
" print(f\" {i+1} (score={score:.2f}): {tokenizer.decode(seq[0], skip_special_tokens=True)}\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "lipogram_private",
"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.14.3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}