{ "cells": [ { "metadata": { "id": "7509f606" }, "cell_type": "markdown", "source": [ "# Assignment 3: Transformer + Machine Translation!\n", "\n", "Welcome to Assignment 3!\n", "\n", "In this exercise, you'll dive into **Neural Machine Translation (NMT)** using **Transformers**, one of the most powerful architectures in natural language processing.\n", "We'll take a bilingual dataset (French → English) and build a **custom Transformer model** from scratch to translate it.\n", "\n", "---\n", "\n", "## Coding Assignment Overview\n", "\n", "The sections below follow the order you will work through them:\n", "\n", "1. **Set up HuggingFace.** Create an account and log in, so you can upload and reload your model checkpoints for convenience (and grading).\n", "2. **Preprocess the dataset.** Encode French–English sentence pair into single token sequences and build the mask that decides which positions the loss is computed on.\n", "3. **Build your own Transformer** model (decoder-only architecture).\n", "4. **Train the model** with a cross-entropy loss (ignoring masked positions), tracking validation loss and BLEU (the BLEU evaluation code is provided for you), then upload your best checkpoint to HuggingFace.\n", "5. **Evaluate** your trained model by reading example translations and computing BLEU on the test set.\n", "6. **Compare** your model with an existing translation model and reflect on what could improve it.\n", "\n", "\n", "## How to Work Through This Notebook\n", "This notebook contains **instructional cells**, **coding tasks**, and **written questions**.\n", "- **Explanation cells** will help you understand key ideas.\n", "- **Coding tasks** are marked with `# BEGIN/END CODING TODO`. Fill in the code between the markers.\n", "- **Written questions** are marked with `BEGIN/END WRITTEN TODO`. Write your answers between the markers.\n", "\n", "Try to **read the comments carefully**, **run cells sequentially**, and **experiment** with hyperparameters or model design choices — exploration is encouraged!\n", "\n", "---\n", "\n", "## Getting Started\n", "\n", "1. **Connect to a GPU.** Training is much faster on a GPU. In Colab, go to **Runtime → Change runtime type**, select **T4 GPU**, and click **Save**. Then click **Connect** in the top right.\n", "2. **Install packages.** Run the cell below. It installs everything the notebook needs and prints whether a GPU is available." ], "id": "1b031224afa18fc1" }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "!pip install -q datasets tokenizers sacrebleu tensorboard sentencepiece sacremoses\n", "# Upgrade these two together so their versions are compatible with each other.\n", "!pip install -q -U transformers huggingface_hub\n", "\n", "import random\n", "import numpy as np\n", "import torch\n", "\n", "# One seed for the whole notebook, so your reported numbers are reproducible.\n", "SEED = 4705\n", "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)\n", "torch.cuda.manual_seed_all(SEED)\n", "\n", "print(\"torch:\", torch.__version__, \"| cuda available:\", torch.cuda.is_available())\n", "if torch.cuda.is_available():\n", " print(\"gpu :\", torch.cuda.get_device_name(0))\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")" ], "id": "ef90823b6e5c92e9" }, { "cell_type": "code", "id": "427c48cd", "metadata": { "id": "427c48cd" }, "source": [ "import math\n", "from tqdm.auto import tqdm\n", "\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torch.utils.data import DataLoader\n", "from torch.optim.lr_scheduler import LambdaLR\n", "from torch.utils.tensorboard import SummaryWriter\n", "\n", "from datasets import load_dataset\n", "\n", "from tokenizers import Tokenizer" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## HuggingFace Setup\n", "\n", "At the end of training, you will upload your best checkpoint to your own HuggingFace repo, so that you can reload it later for more convenient evaluation.\n", "\n", "If you don't have an account yet:\n", "1) Create a HuggingFace account with your student email: https://huggingface.co/join\n", "\n", "2) Run the cell below and follow the instruction to login. Select ```Do not grant Hugging Face CLI access to any of your organizations``` unless you plan to upload your checkpoint to an organization. By default, it is uploaded to your personal account (`HF_USERNAME`), which does not need organization access." ], "id": "af54ee5882443696" }, { "cell_type": "code", "metadata": {}, "source": [ "from huggingface_hub import notebook_login\n", "notebook_login()\n", "\n", "# --------------------------------- BEGIN CODING TODO\n", "# Set your HuggingFace username, which will be used for uploading your checkpoint\n", "HF_USERNAME = \"YOUR_HUGGINGFACE_USERNAME\"\n", "# --------------------------------- END CODING TODO\n", "\n", "assert HF_USERNAME != \"YOUR_HUGGINGFACE_USERNAME\", \"Set HF_USERNAME to your HuggingFace username above, then re-run this cell.\"\n", "\n", "# This is where your checkpoint will be uploaded to. Do not change it, so we can find it for grading.\n", "CKPT_PATH = f\"{HF_USERNAME}/coms4705-a3-mt-transformer\"" ], "execution_count": null, "outputs": [], "id": "39546f9a7147a512" }, { "cell_type": "markdown", "id": "7b9181b0", "metadata": { "id": "7b9181b0" }, "source": [ "## Preprocess the French–English Dataset\n", "\n", "In this section, we will turn raw French–English pairs from the **Multi30k** dataset into batches the model can train on\n", "\n", "We'll use a **Byte-Pair Encoding (BPE)** tokenizer that was pretrained on a large French–English dataset (WMT14, ~1M sentence pairs).\n", "This tokenizer already includes special tokens like `` and ``.\n", "\n", "By the end of this section, the DataLoaders will give you batches that are already tokenized, padded, shifted and masked." ] }, { "cell_type": "code", "id": "ecf97eef", "metadata": { "id": "ecf97eef" }, "source": [ "# View the dataset\n", "dataset = load_dataset(\"coms4705-hewitt/multi30k_fr-en\")\n", "print(dataset)\n", "\n", "train_dataset = dataset[\"train\"]\n", "val_dataset = dataset[\"validation\"]\n", "test_dataset = dataset[\"test\"]\n", "\n", "print(f\"\\nFinal dataset sizes:\")\n", "print(f\"- Train: {len(train_dataset)} examples\")\n", "print(f\"- Validation: {len(val_dataset)} examples\")\n", "print(f\"- Test: {len(test_dataset)} examples\")" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "0cb36472", "vscode": { "languageId": "ruby" } }, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# Get the tokenizer and its special token IDs\n", "tokenizer = Tokenizer.from_pretrained(\"coms4705-hewitt/a2_tokenizer\")\n", "\n", "start_token_id = tokenizer.token_to_id(\"\")\n", "end_token_id = tokenizer.token_to_id(\"\")\n", "pad_token_id = tokenizer.token_to_id(\"\")\n", "\n", "print(\"Special token IDs:\")\n", "print(f\": {start_token_id}\")\n", "print(f\": {end_token_id}\")\n", "print(f\": {pad_token_id}\")" ], "id": "57926bf7a983ca06" }, { "metadata": {}, "cell_type": "markdown", "source": [ "### Pack French-English pair into single masked token sequence\n", "Each French-English sentence pair will be packed into **one** token sequence, so a decoder-only model can generate the\n", "English translation conditioned on the French input:\n", "\n", "```\n", "[ french tokens ] [ english tokens ] \n", "```\n", "\n", "`` marks where the French ends and the translation begins.\n", "\n", "`encode_example` below builds that sequence. Run it to see the packing on a real pair." ], "id": "2b0cea16adeec2a6" }, { "metadata": { "id": "6426190f" }, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "def encode_example(example, tokenizer, verbose=False):\n", " \"\"\"Tokenize one French-English pair into a single sequence.\n", "\n", " Set verbose=True to print each token id under the string it stands for.\n", " \"\"\"\n", " fr = example[\"translation\"][\"fr\"]\n", " en = example[\"translation\"][\"en\"]\n", " enc = tokenizer.encode(fr, en) # tokenizer automatically adds and tokens\n", "\n", " if verbose:\n", " # A word the vocabulary doesn't cover whole gets split into pieces\n", " # (e.g. \"buissons\" -> \"bu\" + \"issons\").\n", " toks = [t.replace(\"\", \"\") for t in enc.tokens]\n", " widths = [max(len(t), len(str(i))) for i, t in zip(enc.ids, toks)]\n", "\n", " print(\"French :\", fr)\n", " print(\"English:\", en)\n", " print(f\"\\n{len(enc.ids)} tokens ( marks where the French ends and the English begins)\\n\")\n", " print(f\"{'tokens':<8}\" + \" \".join(f\"{t:<{w}}\" for t, w in zip(toks, widths)))\n", " print(f\"{'ids':<8}\" + \" \".join(f\"{i:<{w}}\" for i, w in zip(enc.ids, widths)))\n", "\n", " return {\"input_ids\": enc.ids}\n", "\n", "\n", "_ = encode_example(train_dataset[0], tokenizer, verbose=True) # Change the idx for train_datset to explore other examples" ], "id": "b4f6d7e18917fa79" }, { "metadata": {}, "cell_type": "markdown", "source": [ "The model predicts a next token at *every* position, including while it is still reading French. But we only care about how well it translates, so some tokens should not contribute to the loss.\n", "\n", "BEGIN WRITTEN TODO\n", "\n", "Here is a single training example:\n", "```\n", "Un chien court . A dog runs . \n", "```\n", "\n", "**Q1.** Which tokens should be and should not be masked out of the loss (i.e., the model is not scored on predicting them)? Why?\n", "\n", "A: \\\n", "\n", "Now suppose a batch has two examples of different lengths. We pad the shorter one with `` tokens so both rows have the same length:\n", "```\n", "Un chien court . A dog runs . \n", "Deux hommes . Two men . \n", "```\n", "\n", "**Q2.** Why do we pad the sequences?\n", "\n", "A: \\\n", "\n", "**Q3.** Which tokens should be masked in each row now? Why?\n", "\n", "A: \\\n", "\n", "END WRITTEN TODO\n", "\n", "Now implement `create_masked_target` below. It receives the batch already shifted for next-token prediction (`inputs = tokens[:, :-1]`, `targets = tokens[:, 1:]`) and should return `targets` with every masked position set to `ignore_index`. Keep in mind that `` can appear at a different position in each row." ], "id": "7af959569dfae70" }, { "cell_type": "code", "id": "d4923905", "metadata": { "id": "d4923905" }, "source": [ "def create_masked_target(inputs, targets, start_token_id, pad_token_id, ignore_index=-1):\n", " # --------------------------------- BEGIN CODING TODO\n", " # Steps to implement:\n", " # 1. For each sequence, find the first token in inputs\n", " # 2. Create mask: True for targets to include in the loss (from the position in inputs onward,\n", " # i.e. the English tokens and )\n", " # 3. Set the French tokens, and tokens in targets to ignore_index\n", " # 4. Return masked target tensor\n", "\n", " # --------------------------------- END CODING TODO" ], "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "MAX_LEN = 79 # determined from dataset max seq len (after tokenization)\n", "IGNORE_INDEX = -1 # positions set to this value contribute nothing to the loss\n", "\n", "def collate_batch(examples):\n", " \"\"\"Turn a list of raw examples into a training-ready batch.\n", "\n", " All preprocessing happens here, once, at data-loading time:\n", " 1. tokenize each pair into `[fr] [en] `\n", " 2. pad to this batch's longest sequence\n", " 3. apply the LM shift, so inputs[j] is used to predict targets[j]\n", " 4. apply create_masked_target, so the loss only covers the English side\n", "\n", " Returns (inputs, masked_targets, key_padding_mask) -- the training and validation\n", " loops use these directly and do no masking of their own.\n", " \"\"\"\n", " tokens_list = [encode_example(ex, tokenizer) for ex in examples]\n", " ids_list = [torch.tensor(toks[\"input_ids\"], dtype=torch.long) for toks in tokens_list]\n", "\n", " # pad only to this batch's max (instead of the model max_len)\n", " # We check for max_len + 1 because we will be splitting the sequence, `[fr] [en] `, for LM prediction\n", " T = max(x.size(0) for x in ids_list)\n", " assert T <= MAX_LEN + 1, f\"Sequence length {T} exceeds model max_len {MAX_LEN}\"\n", "\n", " B = len(ids_list)\n", "\n", " token_ids = torch.full((B, T), pad_token_id, dtype=torch.long)\n", " for i, seq in enumerate(ids_list):\n", " seq = seq[:T]\n", " token_ids[i, :seq.size(0)] = seq\n", "\n", " # LM shift: predict next token\n", " inputs = token_ids[:, :-1] # [B, T-1]\n", " targets = token_ids[:, 1:] # [B, T-1]\n", " key_padding_mask = (inputs == pad_token_id) # True where PAD\n", "\n", " # Mask French tokens, and so loss is computed on English only.\n", " # This calls YOUR create_masked_target from the cell above.\n", " masked_targets = create_masked_target(\n", " inputs, targets, start_token_id, pad_token_id, ignore_index=IGNORE_INDEX\n", " )\n", "\n", " return inputs, masked_targets, key_padding_mask" ], "id": "89ed9c05457d1dd1", "outputs": [], "execution_count": null }, { "cell_type": "code", "metadata": {}, "source": [ "BATCH_SIZE = 128\n", "\n", "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True, collate_fn=collate_batch)\n", "val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, pin_memory=True, collate_fn=collate_batch)\n", "test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, pin_memory=True, collate_fn=collate_batch)\n", "\n", "print(f\"Train loader: {len(train_loader):>3} batches ({len(train_dataset)} examples)\")\n", "print(f\"Val loader: {len(val_loader):>3} batches ({len(val_dataset)} examples)\")\n", "print(f\"Test loader: {len(test_loader):>3} batches ({len(test_dataset)} examples)\")\n", "\n", "# Pull one batch to check the preprocessing pipeline end to end.\n", "try:\n", " _inputs, _masked_targets, _kpm = next(iter(train_loader))\n", " scored = (_masked_targets != IGNORE_INDEX).sum().item()\n", " print(\"\\nOne batch:\")\n", " print(f\" inputs {tuple(_inputs.shape)}\")\n", " print(f\" masked_targets {tuple(_masked_targets.shape)}\")\n", " print(f\" key_padding_mask {tuple(_kpm.shape)}\")\n", " print(f\" positions scored: {scored} / {_masked_targets.numel()} \"\n", " f\"({100 * scored / _masked_targets.numel():.1f}% -- the English side)\")\n", "except (TypeError, AttributeError, NameError):\n", " print(\"\\n(Could not build a batch yet. Implement create_masked_target above, \"\n", " \"then re-run this cell.)\")" ], "id": "6f98ac7516c7af66", "outputs": [], "execution_count": null }, { "cell_type": "markdown", "id": "d658794f", "metadata": { "id": "d658794f" }, "source": [ "# Transformer\n", "Now we will start building the decoder-only Transformer model!" ] }, { "cell_type": "code", "id": "c545f75b", "metadata": { "id": "c545f75b" }, "source": [ "class MultiHeadAttention(nn.Module):\n", " def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1, block_size: int = 1024):\n", " super().__init__()\n", " assert d_model % num_heads == 0\n", " self.h = num_heads\n", " self.d_model = d_model\n", " self.dk = d_model // num_heads\n", " self.qkv = nn.Linear(d_model, 3 * d_model)\n", " self.o = nn.Linear(d_model, d_model)\n", " self.attn_drop = nn.Dropout(dropout)\n", " self.resid_drop = nn.Dropout(dropout)\n", " self.register_buffer(\"causal_mask\", torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size))\n", "\n", " def forward(self, x, key_padding_mask: torch.Tensor | None = None):\n", " B, T, D = x.size()\n", " q, k, v = self.qkv(x).split(self.d_model, dim=2)\n", " k = k.view(B, T, self.h, self.dk).transpose(1, 2)\n", " q = q.view(B, T, self.h, self.dk).transpose(1, 2)\n", " v = v.view(B, T, self.h, self.dk).transpose(1, 2)\n", "\n", " # --------------------------------- BEGIN CODING TODO\n", " att = # compute scaled dot-product attention\n", " att = # apply causal mask to prevent attending to future tokens\n", " att = # if key_padding_mask is provided, mask out padded positions\n", " att = # apply softmax over attention scores\n", " att = # apply dropout to attention weights\n", " y = # compute weighted sum\n", " # --------------------------------- END CODING TODO\n", "\n", " y = y.transpose(1, 2).contiguous().view(B, T, D)\n", " y = self.resid_drop(self.o(y))\n", " return y\n", "\n", "class FeedForward(nn.Module):\n", " def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):\n", " super().__init__()\n", " self.net = nn.Sequential(\n", " nn.Linear(d_model, d_ff),\n", " nn.GELU(),\n", " nn.Linear(d_ff, d_model),\n", " nn.Dropout(dropout),\n", " )\n", " def forward(self, x):\n", " # --------------------------------- BEGIN CODING TODO\n", " x = # pass input through the feed-forward network (2 linear layers + GELU)\n", " # --------------------------------- END CODING TODO\n", " return x\n", "\n", "class Block(nn.Module):\n", " def __init__(self, d_model, num_heads, d_ff, dropout=0.1, block_size=1024):\n", " super().__init__()\n", " self.ln1 = nn.LayerNorm(d_model)\n", " self.mha = MultiHeadAttention(d_model, num_heads, dropout, block_size)\n", " self.ln2 = nn.LayerNorm(d_model)\n", " self.mlp = FeedForward(d_model, d_ff, dropout)\n", "\n", " def forward(self, x, key_padding_mask: torch.Tensor | None = None):\n", "\n", " # --------------------------------- BEGIN CODING TODO\n", " x = # apply LayerNorm + MultiHeadAttention + residual connection\n", " x = # apply LayerNorm + FeedForward + residual connection\n", " # --------------------------------- END CODING TODO\n", "\n", " return x\n", "\n", "class DecoderOnlyTransformer(nn.Module):\n", " def __init__(self, vocab_size, d_model=256, n_layers=4, num_heads=4, d_ff=512, max_len=1024, dropout=0.1):\n", " super().__init__()\n", " self.n_layers = n_layers # Store for GPT-style weight init\n", " self.max_len = max_len # Store max_len for generation\n", " self.tok = nn.Embedding(vocab_size, d_model)\n", " self.pos = nn.Embedding(max_len, d_model) # Learned positional embeddings (like GPT)\n", " self.drop = nn.Dropout(dropout)\n", " self.blocks = nn.ModuleList([Block(d_model, num_heads, d_ff, dropout, max_len) for _ in range(n_layers)])\n", " self.ln = nn.LayerNorm(d_model)\n", " self.head = nn.Linear(d_model, vocab_size, bias=False)\n", "\n", " # Weight tying: tie the weights of the token embedding and output head\n", " self.head.weight = self.tok.weight\n", "\n", " # Apply GPT-style weight initialization\n", " self.apply(self._init_weights)\n", " # Special scaled initialization for all projection layers (output projections in residual paths)\n", " for pn, p in self.named_parameters():\n", " # Apply to attention output projection (mha.o) and MLP output projection (mlp.net.2)\n", " if pn.endswith('mha.o.weight') or pn.endswith('mlp.net.2.weight'):\n", " torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * n_layers))\n", "\n", " def _init_weights(self, module):\n", " \"\"\"GPT-style weight initialization for better training stability\"\"\"\n", " if isinstance(module, nn.Linear):\n", " torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n", " if module.bias is not None:\n", " torch.nn.init.zeros_(module.bias)\n", " elif isinstance(module, nn.Embedding):\n", " torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n", " elif isinstance(module, nn.LayerNorm):\n", " torch.nn.init.zeros_(module.bias)\n", " torch.nn.init.ones_(module.weight)\n", "\n", " def configure_optimizers(self, learning_rate, weight_decay, betas):\n", " \"\"\"\n", " GPT-style optimizer configuration that separates parameters into decay/no_decay groups.\n", " Weight decay is only applied to weight matrices in Linear layers, not biases or LayerNorm.\n", " Handles weight tying between token embedding and output head.\n", " \"\"\"\n", " # separate out all parameters to those that will and won't experience regularizing weight decay\n", " decay = set()\n", " no_decay = set()\n", " whitelist_weight_modules = (torch.nn.Linear, )\n", " blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n", "\n", " # Get actual parameters that exist (handles weight tying)\n", " param_dict = {pn: p for pn, p in self.named_parameters()}\n", "\n", " for mn, m in self.named_modules():\n", " for pn, p in m.named_parameters():\n", " fpn = '%s.%s' % (mn, pn) if mn else pn # full param name\n", "\n", " # Skip if parameter doesn't exist (due to weight tying)\n", " if fpn not in param_dict:\n", " continue\n", "\n", " if pn.endswith('bias'):\n", " # all biases will not be decayed\n", " no_decay.add(fpn)\n", " elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):\n", " # weights of whitelist modules will be weight decayed\n", " decay.add(fpn)\n", " elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):\n", " # weights of blacklist modules will NOT be weight decayed\n", " no_decay.add(fpn)\n", "\n", " # validate that we considered every parameter\n", " inter_params = decay & no_decay\n", " union_params = decay | no_decay\n", " assert len(inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n", " assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n", " % (str(param_dict.keys() - union_params), )\n", "\n", " # create the pytorch optimizer object\n", " optim_groups = [\n", " {\"params\": [param_dict[pn] for pn in sorted(list(decay))], \"weight_decay\": weight_decay},\n", " {\"params\": [param_dict[pn] for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n", " ]\n", " optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)\n", " return optimizer\n", "\n", " def forward(self, x, key_padding_mask: torch.Tensor | None = None):\n", " B, T = x.shape\n", " device = x.device\n", "\n", " # --------------------------------- BEGIN CODING TODO\n", " pos = # create position indices from 0..T-1 and expand for batch\n", " tok_emb = # compute token embeddings\n", " pos_emb = # compute positional embeddings\n", " x_emb = # sum token and position embeddings, then apply dropout\n", " y = # pass embeddings sequentially through all transformer blocks\n", " y = # apply final LayerNorm\n", " logits = # compute logits via final linear head\n", " # --------------------------------- END CODING TODO\n", "\n", " return logits\n", "\n", " @torch.no_grad()\n", " def generate(self, x, max_new_tokens, eos_id=None):\n", " for _ in range(max_new_tokens):\n", " # Stop if sequence length would exceed max_len\n", " if x.size(1) >= self.max_len:\n", " break\n", "\n", " logits = self.forward(x)[:, -1, :]\n", " next_id = torch.argmax(logits, dim=-1, keepdim=True)\n", " x = torch.cat([x, next_id], dim=1)\n", " if eos_id is not None and (next_id == eos_id).all():\n", " break\n", " return x" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Evaluation Metric\n", "\n", "We need a way to score a translation. `evaluate_bleu` below generates a translation for each French sentence and compares it against the reference English, scoring the result with **corpus BLEU**." ], "id": "a80680a2b4183b84" }, { "cell_type": "code", "id": "6035b047", "metadata": { "id": "6035b047" }, "source": [ "# Import BLEU evaluation\n", "from sacrebleu import corpus_bleu\n", "\n", "# Decoding setting used everywhere we generate or score translations (validation during training,\n", "# the examples, the final BLEU and the OpusMT comparison): greedy, up to MAX_NEW_TOKENS new tokens.\n", "MAX_NEW_TOKENS = 50\n", "\n", "def evaluate_bleu(model, dataset, tokenizer, device, max_new_tokens, num_samples):\n", " \"\"\"\n", " Evaluate BLEU score on a subset of the dataset, generating up to `max_new_tokens` tokens per translation.\n", " `num_samples` is how many random examples to score; pass None to score the whole dataset.\n", " \"\"\"\n", " model.eval()\n", " predictions = []\n", " references = []\n", "\n", " # Sample a subset for BLEU evaluation (to save time during training)\n", " if num_samples is not None:\n", " sample_indices = torch.randperm(len(dataset))[:num_samples]\n", " else:\n", " sample_indices = list(range(len(dataset)))\n", "\n", " with torch.no_grad():\n", " for idx in sample_indices:\n", " ex = dataset[int(idx)]\n", " fr_text = ex[\"translation\"][\"fr\"]\n", " en_text = ex[\"translation\"][\"en\"]\n", "\n", " # Encode French text and add start token\n", " fr_enc = tokenizer.encode(fr_text).ids\n", " fr_enc += [start_token_id]\n", "\n", " # Generate translation\n", " out = model.generate(\n", " torch.tensor([fr_enc], device=device),\n", " max_new_tokens=max_new_tokens,\n", " eos_id=end_token_id\n", " )\n", "\n", " # --------------------------------- BEGIN CODING TODO\n", " # Steps to implement:\n", " # 1. Keep only the newly generated token IDs, i.e. everything after the prompt (French + )\n", " # 2. If the model produced , cut the sequence there ( is not part of the translation)\n", " # 3. Decode the remaining IDs into text and append it to `predictions`\n", " # 4. Append the reference English text `en_text` to `references`\n", "\n", " # --------------------------------- END CODING TODO\n", "\n", "\n", " # Calculate corpus BLEU over all (prediction, reference) pairs.\n", " # Tip: BLEU is not the main focus of this assignment, but we encourage you to read how\n", " # sacrebleu computes it: https://github.com/mjpost/sacrebleu\n", " if not predictions:\n", " return 0.0\n", "\n", " # Our tokenizer decodes punctuation as separate tokens (\"a dog runs .\"), which makes sacrebleu warn\n", " # that the output looks tokenized. force=True silences the warning.\n", " return corpus_bleu(predictions, [references], force=True).score" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "id": "51af0ee9", "metadata": { "id": "51af0ee9" }, "source": [ "# Training\n", "\n", "We will use tensorboard to visualize our training process." ] }, { "cell_type": "code", "id": "a25bae33", "metadata": { "id": "a25bae33" }, "source": [ "# Model config. These settings reached ~1.4 validation loss and >30 BLEU on val and test.\n", "VOCAB_SIZE = tokenizer.get_vocab_size()\n", "D_MODEL = 128\n", "N_LAYERS = 4\n", "NUM_HEADS = 4\n", "D_FF = 512\n", "\n", "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "model = DecoderOnlyTransformer(\n", " vocab_size=VOCAB_SIZE,\n", " d_model=D_MODEL,\n", " n_layers=N_LAYERS,\n", " num_heads=NUM_HEADS,\n", " d_ff=D_FF,\n", " max_len=MAX_LEN,\n", ").to(device)" ], "outputs": [], "execution_count": null }, { "metadata": { "id": "58dc7530" }, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "%load_ext tensorboard\n", "%tensorboard --logdir runs --reload_interval 1" ], "id": "7f4c2308c24a746e" }, { "cell_type": "code", "metadata": {}, "source": [ "# ---------------------------- Training Loop ----------------------------\n", "# - Trains for `num_epochs` with cross-entropy loss. Batches arrive already masked by\n", "# `collate_batch`, so French tokens, and contribute nothing to the loss.\n", "# - AdamW with inverse square root warmup, plus gradient clipping for stability.\n", "# - After each epoch: validation loss + BLEU, and the best model (by val loss) is saved.\n", "# - TensorBoard logs training loss, validation loss, BLEU and learning rate.\n", "#\n", "# Tip: experiment with the hyperparameters below. You can scroll up and see their effect in TensorBoard above.\n", "# -----------------------------------------------------------------------\n", "\n", "import os\n", "from datetime import datetime\n", "\n", "# ---- Hyperparameters ----\n", "num_epochs = 30\n", "base_lr = 1e-3\n", "weight_decay = 0.1\n", "betas = (0.9, 0.95)\n", "grad_norm_clip = 1.0\n", "warmup_steps = 500\n", "val_bleu_samples = 50 # random validation examples scored for BLEU after each epoch (full set is slow)\n", "\n", "# ---- Optimizer, LR schedule and loss ----\n", "optimizer = model.configure_optimizers(base_lr, weight_decay, betas)\n", "\n", "def inv_sqrt_lambda(step):\n", " \"\"\"Linear warmup for `warmup_steps`, then decay proportional to 1/sqrt(step).\"\"\"\n", " step = max(1, step)\n", " return (warmup_steps ** 0.5) * min(step ** -0.5, step / (warmup_steps ** 1.5))\n", "\n", "scheduler = LambdaLR(optimizer, lr_lambda=inv_sqrt_lambda)\n", "criterion = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)\n", "\n", "total_steps = len(train_loader) * num_epochs\n", "print(f\"Total steps: {total_steps}, Warmup steps: {warmup_steps}\")\n", "\n", "# ---- Logging and checkpointing ----\n", "timestamp = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n", "experiment_name = f\"dm={D_MODEL}_multi30k_{timestamp}\"\n", "log_dir = f\"runs/{experiment_name}\"\n", "writer = SummaryWriter(log_dir)\n", "print(f\"TensorBoard logs will be saved to: {log_dir}\")\n", "\n", "os.makedirs(\"ckpts\", exist_ok=True)\n", "best_model_path = f\"ckpts/{experiment_name}.pt\"\n", "best_val_loss = float(\"inf\")\n", "\n", "\n", "def compute_loss(model, inputs, masked_targets, key_padding_mask):\n", " inputs = inputs.to(device)\n", " masked_targets = masked_targets.to(device)\n", " kpm = key_padding_mask.to(device)\n", " logits = model(inputs, key_padding_mask=kpm)\n", " return criterion(logits.reshape(-1, VOCAB_SIZE), masked_targets.reshape(-1))\n", "\n", "\n", "@torch.no_grad()\n", "def evaluate_loss(model, loader):\n", " model.eval()\n", " total = sum(compute_loss(model, *batch).item() for batch in loader)\n", " model.train()\n", " return total / len(loader)\n", "\n", "\n", "def save_checkpoint(path, epoch, val_loss):\n", " torch.save({\n", " \"epoch\": epoch,\n", " \"model_state_dict\": model.state_dict(),\n", " \"optimizer_state_dict\": optimizer.state_dict(),\n", " \"val_loss\": val_loss,\n", " \"vocab_size\": VOCAB_SIZE,\n", " \"model_config\": {\n", " \"d_model\": D_MODEL,\n", " \"n_layers\": N_LAYERS,\n", " \"num_heads\": NUM_HEADS,\n", " \"d_ff\": D_FF,\n", " \"max_len\": MAX_LEN,\n", " },\n", " }, path)\n", "\n", "\n", "# ---- Training loop ----\n", "global_step = 0\n", "model.train()\n", "for epoch in range(1, num_epochs + 1):\n", " running_loss = 0.0\n", " for i, batch in enumerate(train_loader, start=1):\n", " loss = compute_loss(model, *batch)\n", " optimizer.zero_grad()\n", " loss.backward()\n", " torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm_clip)\n", " optimizer.step()\n", " scheduler.step()\n", "\n", " last_lr = scheduler.get_last_lr()[0]\n", " writer.add_scalar(\"Loss/Train_Batch\", loss.item(), global_step)\n", " writer.add_scalar(\"Learning_Rate\", last_lr, global_step)\n", " global_step += 1\n", "\n", " running_loss += loss.item()\n", " if i % 100 == 0:\n", " print(f\"Epoch {epoch}, Step {i}, Avg Loss: {running_loss / 100:.4f}, Learning Rate: {last_lr:.6f}\")\n", " running_loss = 0.0\n", "\n", " # ---- Validation (end of each epoch) ----\n", " val_loss = evaluate_loss(model, val_loader)\n", " val_bleu = evaluate_bleu(model, val_dataset, tokenizer, device, MAX_NEW_TOKENS, num_samples=val_bleu_samples)\n", " model.train() # evaluate_bleu switches the model to eval mode\n", "\n", " writer.add_scalar(\"Loss/Validation\", val_loss, global_step)\n", " writer.add_scalar(\"BLEU/Validation\", val_bleu, global_step)\n", " print(f\"*** Epoch {epoch}, Validation Loss: {val_loss:.4f}, BLEU: {val_bleu:.2f} ***\")\n", "\n", " if val_loss < best_val_loss:\n", " best_val_loss = val_loss\n", " save_checkpoint(best_model_path, epoch, val_loss)\n", " print(f\"*** New best model saved! Val Loss: {val_loss:.4f} ***\")\n", "\n", "writer.close()\n", "print(f\"\\nTraining completed! Best validation loss: {best_val_loss:.4f}\")\n", "print(f\"Best model saved as: {best_model_path}\")\n", "print(f\"TensorBoard logs saved in: {log_dir}\")" ], "id": "753d8bf1d534cf47", "outputs": [], "execution_count": null }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "# ---------------------------- Upload Checkpoint ----------------------------\n", "# Upload your best checkpoint to HuggingFace. This is what we use for grading, so make sure it succeeds.\n", "# ---------------------------------------------------------------------------\n", "from huggingface_hub import HfApi\n", "\n", "api = HfApi()\n", "api.create_repo(CKPT_PATH, exist_ok=True)\n", "api.upload_file(\n", " path_or_fileobj=best_model_path,\n", " path_in_repo=\"best_model.pt\",\n", " repo_id=CKPT_PATH,\n", ")\n", "print(f\"Uploaded {best_model_path} to https://huggingface.co/{CKPT_PATH}\")" ], "id": "4b5b888e" }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "def load_best_model(model_path, device):\n", " \"\"\"Rebuild the model from a saved checkpoint and restore its weights.\"\"\"\n", " checkpoint = torch.load(model_path, map_location=device)\n", " model = DecoderOnlyTransformer(\n", " vocab_size=checkpoint[\"vocab_size\"], **checkpoint[\"model_config\"]\n", " ).to(device)\n", " model.load_state_dict(checkpoint[\"model_state_dict\"])\n", " print(f\"Loaded best model from epoch {checkpoint['epoch']} (val_loss: {checkpoint['val_loss']:.4f})\")\n", " return model, checkpoint" ], "id": "548394ffc63182a1" }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# ---------------------------- Load Checkpoint ----------------------------\n", "# Load your best checkpoint as `model` for generating translations and evaluation.\n", "# By default this loads the local checkpoint from training. To load the one you uploaded\n", "# instead (e.g. after your Colab session restarted), uncomment the two lines below.\n", "# -------------------------------------------------------------------------\n", "# from huggingface_hub import hf_hub_download\n", "# best_model_path = hf_hub_download(CKPT_PATH, \"best_model.pt\")\n", "\n", "model, checkpoint_info = load_best_model(best_model_path, device)" ], "id": "75180a52fe1a83d" }, { "metadata": {}, "cell_type": "markdown", "source": [ "## Evaluating the Trained Model\n", "\n", "Let's see how your trained model does on the test set. First, read a few example translations next to their references and note which sentences it handles well and which it struggles with.\n", "Then we reuse `evaluate_bleu` from the Evaluation Metric section to compute BLEU on the full test set, with the same decoding settings (`MAX_NEW_TOKENS`).\n", "\n", "Tips:\n", "- Try changing `max_new_tokens` to generate longer or shorter translations.\n", "- Observe which sentences the model translates well and which it struggles with.\n", "- You can modify the loop to test your own French sentences." ], "id": "88a278fdc6536c50" }, { "metadata": {}, "cell_type": "code", "outputs": [], "execution_count": null, "source": [ "model.eval()\n", "for i, ex in enumerate(test_dataset):\n", " fr_enc = tokenizer.encode(ex[\"translation\"][\"fr\"]).ids\n", " fr_enc += [start_token_id]\n", " out = model.generate(torch.tensor([fr_enc], device=device), max_new_tokens=MAX_NEW_TOKENS, eos_id=end_token_id)\n", " print(f\"Output: {tokenizer.decode(out[0].tolist())}\")\n", " print(f\"REFERENCE translation: {ex['translation']['en']}\")\n", " print()\n", " if i == 10:\n", " break" ], "id": "426ca2989eb78419" }, { "cell_type": "code", "id": "8993c76e", "metadata": { "id": "8993c76e" }, "source": [ "test_bleu = evaluate_bleu(model, test_dataset, tokenizer, device, MAX_NEW_TOKENS, num_samples=None)\n", "print(\"Test BLEU:\", test_bleu)" ], "outputs": [], "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Comparing with an Existing Model\n", "\n", "How good is your model's BLEU? To put it in context, we compare it with [Helsinki-NLP/opus-mt-fr-en](https://huggingface.co/Helsinki-NLP/opus-mt-fr-en) as our baseline, an existing French-to-English translation model. It is an encoder-decoder Transformer trained on a much larger amount of French–English data.\n", "\n", "The cell below translates the same test set with OpusMT, using the same decoding settings as your model (greedy, up to `MAX_NEW_TOKENS` new tokens), and prints both models' test BLEU." ], "id": "fa9736c502f38fc9" }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "from tqdm.auto import trange\n", "from huggingface_hub import snapshot_download\n", "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM\n", "\n", "# Download the model files first, then load from the local folder.\n", "opus_path = snapshot_download(\"Helsinki-NLP/opus-mt-fr-en\")\n", "opus_tokenizer = AutoTokenizer.from_pretrained(opus_path)\n", "opus_model = AutoModelForSeq2SeqLM.from_pretrained(opus_path).to(device).eval()\n", "\n", "\n", "@torch.no_grad()\n", "def generate_translations(dataset, batch_size=64):\n", " \"\"\"Translate `dataset` with OpusMT in batches. Returns (predictions, references).\"\"\"\n", " preds = []\n", " refs = []\n", " for i in trange(0, len(dataset), batch_size):\n", " batch = dataset[i:i + batch_size]\n", " fr_texts = [ex[\"fr\"] for ex in batch[\"translation\"]]\n", " en_texts = [ex[\"en\"] for ex in batch[\"translation\"]]\n", " inputs = opus_tokenizer(fr_texts, return_tensors=\"pt\", padding=True, truncation=True).to(opus_model.device)\n", " # same decoding settings as your model: greedy, up to MAX_NEW_TOKENS new tokens\n", " outputs = opus_model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, num_beams=1)\n", " preds.extend(opus_tokenizer.batch_decode(outputs, skip_special_tokens=True))\n", " refs.extend(en_texts)\n", " return preds, refs\n", "\n", "\n", "opus_preds, opus_refs = generate_translations(test_dataset)\n", "opus_bleu = corpus_bleu(opus_preds, [opus_refs]).score\n", "\n", "print(f\"Your model test BLEU: {test_bleu:.2f}\")\n", "print(f\"OpusMT test BLEU: {opus_bleu:.2f}\")" ], "id": "1c585eecbfba6153" }, { "cell_type": "markdown", "metadata": {}, "source": [ "BEGIN WRITTEN TODO\n", "\n", "**Q.** How does your model's BLEU compare to OpusMT's? What could be tried to improve your model? Consider anything: the data, the model, the training setup, or how translations are generated. Pick a few ideas and explain why you expect them to help.\n", "\n", "A: \\\n", "\n", "END WRITTEN TODO" ], "id": "4c00e7d117b1d5c9" } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "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.10" } }, "nbformat": 4, "nbformat_minor": 5 }