{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f1a563bc",
   "metadata": {},
   "source": [
    "# Dot Product Attention\n",
    "\n",
    "Time to implement a transformer!\n",
    "\n",
    "Note that these models are getting big, so unless you have a GPU with at least 8 GB memory, it is probably better to work on the transformer notebooks online, for example on [Google Colab](https://colab.research.google.com)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e02ad6",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## Imports\n",
    "\n",
    "Import `datasets`, `os`, `torch`, and `torch.nn.functional` as `F`.\n",
    "\n",
    "Moreover, import the following:\n",
    "1. The function `get_seed`, that you wrote in Notebook 0625,\n",
    "1. The functions `get_binary_accuracy` and `get_binary_cross_entropy`, that you wrote in Notebook 0701,\n",
    "1. The class `Linear` and the function `get_mlp`, that you wrote in Notebook 0706,\n",
    "1. The class `AdamW`, that you wrote in Notebook 0708,\n",
    "1. The class `RMSNorm`, that you wrote in Notebook 0720, and\n",
    "1. The classes `Embedding` and `MeanPool`, and the functions `evaluate_model`, `get_dataloader_random_reshuffle`, `tokenize_corpus`, and `train_supervised`, that you wrote in Notebook 0721."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b28c96b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "import os\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from util_0722 import (\n",
    "    AdamW,\n",
    "    Embedding,\n",
    "    evaluate_model,\n",
    "    get_binary_accuracy,\n",
    "    get_binary_cross_entropy,\n",
    "    get_dataloader_random_reshuffle,\n",
    "    get_mlp,\n",
    "    get_seed,\n",
    "    Linear,\n",
    "    MeanPool,\n",
    "    RMSNorm,\n",
    "    tokenize_corpus,\n",
    "    train_supervised,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06d03d08",
   "metadata": {},
   "source": [
    "### Configuration\n",
    "\n",
    "This will be about the same as in Notebook 0721, besides making sizes smaller so that we may fit in memory.\n",
    "\n",
    "Create a configuration dictionary with the following keys:\n",
    "- `\"dataset_path\"`: `str`  \n",
    "    Make this the ID of the IMDB dataset [1], to be found at https://huggingface.co/datasets/stanfordnlp/imdb.\n",
    "- `\"dataset_preprocessed_path\"` : `str`  \n",
    "    Below, I'm going to suggest saving the preprocessed dataset at this location. It should have a `pt` extension.\n",
    "- `\"device\"`: `torch.device | int | str`  \n",
    "    The device identifier.\n",
    "- `\"ensemble_shape\"`: `tuple[int]`  \n",
    "    On my 8GB GPU, I had to make this `(3,)`.\n",
    "- `\"hyperparameter_raw_init_distributions\"`, `\"hyperparameter_raw_perturb\"`, `\"hyperparameter_transforms\"` : `dict`  \n",
    "    These three dictionaries are going to determine how the hyperparameters are tuned. We'll tune the following hyperparameters:\n",
    "    1. Epsilon $\\epsilon$.\n",
    "    2. Learning rate $\\eta$.\n",
    "    3. Weight decay $\\lambda$.\n",
    "    4. First moment moving average decay rate $\\beta_1$.\n",
    "    5. Second moment moving average decay rate $\\beta_2$.\n",
    "\n",
    "    Of these, we don't know the required order of magnitude of the first three. Thus it may be good to make them distributed along $10^\\mathscr D$ where $\\mathscr D$ is a normal or uniform distribution. You can try to center the distributions at the recommended values.\n",
    "\n",
    "    We know that the recommended values of the fourth and fifth are $0.9$ and $0.999$. So it may be best to give them a distribution of the form $1-10^\\mathscr D$.\n",
    "- `\"improvement_threshold:`: `float`  \n",
    "    Make this `1e-4`.\n",
    "- `\"minibatch_size\"`: `int`  \n",
    "    On my 8GB GPU, I had to make this `8`.\n",
    "- `\"minibatch_size_eval\"`: `int`  \n",
    "    On my 8GB GPU, I had to make this `8`.\n",
    "- `\"pbt\"` : `bool`  \n",
    "    Make this `True`.\n",
    "- `\"seed\"`: `int`  \n",
    "    This is for reproducible experiments. Insert any integer.\n",
    "- `\"steps_num\"`: `int`  \n",
    "    Make this `10_001`.\n",
    "- `\"steps_without_improvement`: `int`  \n",
    "    Make this `10_000`.\n",
    "- `\"valid_interval\"`: `int`  \n",
    "    Make this `1000`.\n",
    "- `\"welch_confidence_level\"`: `float`  \n",
    "    We will exploit based on a one-sided Welch $t$-test with this confidence level. Maybe you can try `.8`. Feel free to try out various values here!\n",
    "- `\"welch_sample_size\"`: `int`  \n",
    "    We will exploit based on a one-sided Welch $t$-test on the last this many validation metrics of the population members. To follow the PBT paper, make this `10`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99292a9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {\n",
    "    \"dataset_path\": \"stanfordnlp/imdb\",\n",
    "    \"dataset_preprocessed_path\": \"data/imdb.pt\",\n",
    "    \"device\": \"cuda\",\n",
    "    \"ensemble_shape\": (3,),\n",
    "    \"hyperparameter_raw_init_distributions\": {\n",
    "        \"epsilon\": torch.distributions.Uniform(\n",
    "            torch.tensor(-10, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(-5, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"first_moment_decay\": torch.distributions.Uniform(\n",
    "            torch.tensor(-3, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(0, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"learning_rate\": torch.distributions.Uniform(\n",
    "            torch.tensor(-5, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(-1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"second_moment_decay\": torch.distributions.Uniform(\n",
    "            torch.tensor(-5, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(-1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"weight_decay\": torch.distributions.Uniform(\n",
    "            torch.tensor(-5, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(-1, device=\"cuda\", dtype=torch.float32)\n",
    "        )\n",
    "    },\n",
    "    \"hyperparameter_raw_perturb\": {\n",
    "        \"epsilon\": torch.distributions.Normal(\n",
    "            torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"first_moment_decay\": torch.distributions.Normal(\n",
    "            torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"learning_rate\": torch.distributions.Normal(\n",
    "            torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"second_moment_decay\": torch.distributions.Normal(\n",
    "            torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "        \"weight_decay\": torch.distributions.Normal(\n",
    "            torch.tensor(0, device=\"cuda\", dtype=torch.float32),\n",
    "            torch.tensor(1, device=\"cuda\", dtype=torch.float32)\n",
    "        ),\n",
    "    },\n",
    "    \"hyperparameter_transforms\": {\n",
    "        \"epsilon\": lambda log10: 10 ** log10,\n",
    "        \"first_moment_decay\": lambda x: (1 - 10 ** x).clamp(0, 1),\n",
    "        \"learning_rate\": lambda log10: 10 ** log10,\n",
    "        \"second_moment_decay\": lambda x: (1 - 10 ** x).clamp(0, 1),\n",
    "        \"weight_decay\": lambda log10: 10 ** log10,\n",
    "    },\n",
    "    \"improvement_threshold\": 1e-4,\n",
    "    \"minibatch_size\": 1 << 3,\n",
    "    \"minibatch_size_eval\": 1 << 3,\n",
    "    \"pbt\": True,\n",
    "    \"seed\": 0,\n",
    "    \"steps_num\": 10_001,\n",
    "    \"steps_without_improvement\": 10_000,\n",
    "    \"valid_interval\": 1000,\n",
    "    \"welch_confidence_level\": .8,\n",
    "    \"welch_sample_size\": 10,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a5a03c1",
   "metadata": {},
   "source": [
    "Set the `torch` pseudo-random number generation seed, as per the configuration dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70e493c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(config[\"seed\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0121fb71",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "\n",
    "Let's do this similarly to Notebook 0721. One catch: If the validation set contains a word that the train set does not, then the model did not have a chance to train the corresponding word vector. As we don't use positional encoding yet, we can first tokenize the train set, and in the validation set, we can set `fixed_vocabulary=True` in `tokenize_corpus`. This will make it so that when tokenizing the validation set, the tokenizer will skip words that are not in the training set.\n",
    "\n",
    "I also suggest first checking if the preprocessed dataset is saved, and if so, just load it. Note that when you're first working on this notebook, you may want to erase the preprocessed dataset as we're changing it in the previous paragraph.\n",
    "\n",
    "Verify by printing dataset sizes and maybe decoding some sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b0935f",
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(config[\"dataset_preprocessed_path\"]):\n",
    "    print(f\"Loading preprocessed dataset from {config['dataset_preprocessed_path']}\")\n",
    "    preprocessed = torch.load(config[\"dataset_preprocessed_path\"])\n",
    "    train_tokenized = preprocessed[\"train\"]\n",
    "    valid_tokenized = preprocessed[\"valid\"]\n",
    "    id2token = preprocessed[\"id2token\"]\n",
    "    token2id = preprocessed[\"token2id\"]\n",
    "else:\n",
    "    print(f\"Preprocessing dataset and saving to {config['dataset_preprocessed_path']}\")\n",
    "\n",
    "    dataset = datasets.load_dataset(config[\"dataset_path\"])\n",
    "    val_test = dataset[\"test\"].train_test_split(\n",
    "        seed=get_seed(),\n",
    "        test_size=.5\n",
    "    )\n",
    "    train, valid = dataset[\"train\"], val_test[\"train\"]\n",
    "\n",
    "    train_token_ids, train_indptr, id2token, token2id = tokenize_corpus(\n",
    "        config,\n",
    "        train[\"text\"],\n",
    "        fixed_vocabulary=False,\n",
    "    )\n",
    "\n",
    "    valid_token_ids, valid_indptr, id2token, token2id = tokenize_corpus(\n",
    "        config,\n",
    "        valid[\"text\"],\n",
    "        fixed_vocabulary=True,\n",
    "        id2token=id2token,\n",
    "        token2id=token2id\n",
    "    )\n",
    "\n",
    "    train_tokenized = {\n",
    "        \"token_ids\": train_token_ids,\n",
    "        \"indptr\": train_indptr,\n",
    "        \"label\": train.with_format(\n",
    "            \"torch\",\n",
    "            device=config[\"device\"]\n",
    "        )[\"label\"][:].to(torch.float32)\n",
    "    }\n",
    "    valid_tokenized = {\n",
    "        \"token_ids\": valid_token_ids,\n",
    "        \"indptr\": valid_indptr,\n",
    "        \"label\": valid.with_format(\n",
    "            \"torch\",\n",
    "            device=config[\"device\"]\n",
    "        )[\"label\"][:].to(torch.float32)\n",
    "    }\n",
    "\n",
    "    torch.save(\n",
    "        {\n",
    "            \"train\": train_tokenized,\n",
    "            \"valid\": valid_tokenized,\n",
    "            \"id2token\": id2token,\n",
    "            \"token2id\": token2id\n",
    "        },\n",
    "        config[\"dataset_preprocessed_path\"]\n",
    "    )\n",
    "\n",
    "print(\"Dataset shapes:\")\n",
    "for d in (train_tokenized, valid_tokenized):\n",
    "    for key, value in d.items():\n",
    "        print(key, value.shape)\n",
    "\n",
    "print(\"Decoding example tokenized text:\")\n",
    "indptr = train_tokenized[\"indptr\"]\n",
    "token_ids = train_tokenized[\"token_ids\"][indptr[1000]:indptr[1001]]\n",
    "print(*(id2token[token_id.cpu().item()] for token_id in token_ids))\n",
    "print(\"Example label:\", train_tokenized[\"label\"][1000].item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef72558c",
   "metadata": {},
   "source": [
    "Make a train dataloader, get a minibatch, print key names and shapes. Also, take a minibatch entry and decode the masked portion for verification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02bdb5e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader_train = get_dataloader_random_reshuffle(\n",
    "    config,\n",
    "    train_tokenized,\n",
    "    dataset_size_key=\"label\"\n",
    ")\n",
    "\n",
    "print(\"Minibatch shapes:\")\n",
    "minibatch = next(dataloader_train)\n",
    "for key, value in minibatch.items():\n",
    "    print(key, value.shape)\n",
    "\n",
    "print(\"Decoding example minibatch tokenized text:\")\n",
    "mask = minibatch[\"mask\"][0,0]\n",
    "token_ids = minibatch[\"token_ids\"][0,0,mask]\n",
    "print(*(id2token[token_id.cpu().item()] for token_id in token_ids))\n",
    "print(\"example label:\", minibatch[\"label\"][0,0].item())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16498186",
   "metadata": {},
   "source": [
    "## Architecture\n",
    "\n",
    "### Attention Block\n",
    "\n",
    "Time to make an attention block!\n",
    "\n",
    "We need four linear layers for the key, output, query, and value parameters. We can make these `(embedding_dim,embedding_dim)` and reshape and transpose output to end with `(attention_head_num, sequence_length, key_dim)` in the `forward` method.\n",
    "\n",
    "To actually perform the dot product attention operation, I recommend using [`F.scaled_dot_product_attention`](https://docs.pytorch.org/docs/2.13/generated/torch.nn.functional.scaled_dot_product_attention.html) as it is an optimized implementation. You can supply the mask at the `attn_mask` keyword argument. Note that it needs a tensor with shape ending with `(sequence_length, sequence_length)` while the mask in the minibatch only has one such dimension. This is because the mask in the minibatch says for which position the entry is valid, while in `attn_mask` you want to say for which query-key pairs the entry is valid. You can use broadcasting and the bitwise and `&` operation.\n",
    "\n",
    "For simplicity, let's not use dropout, and use pre-normalization. Note that the `Linear` and `RMSNorm` layers we previously implemented use plain feature tensors as input and output.\n",
    "\n",
    "Don't forget to use a residual connection!\n",
    "\n",
    "For testing, make an embedding, and an attention block, feed them the minibatch, and check output keys and shapes. You can keep sizes low, eg. embedding dimension 16 and 2 heads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b03c85a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionBlock(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Ensemble-ready dot product self attention block. For now, it has\n",
    "    pre-normalization and no positional encoding. We let key and value dimensions be `embedding_dim // attention_head_num`.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    attention_head_num : `int`\n",
    "        The number of attention heads.\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        `\"device\"` : `str`\n",
    "            The device to store parameters on.\n",
    "        `\"ensemble_shape\"` : `tuple[int]`\n",
    "            The shape of the ensemble of affine transformations\n",
    "            the model represents.\n",
    "    embedding_dim : `int`\n",
    "        The number of embedding dimensions.\n",
    "\n",
    "    Calling\n",
    "    -------\n",
    "    The expected input is a dictionary of the following\n",
    "    key-tensor pairs:\n",
    "    `\"features\"`\n",
    "        The feature tensor, of shape\n",
    "        `(..., sequence_dim, feature_dim)`\n",
    "    `\"mask\"`\n",
    "        This mask signals which entry is not a padding element.\n",
    "        It should have size `(..., sequence_dim)`.\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        attention_head_num: int,\n",
    "        config: dict,\n",
    "        embedding_dim: int,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "\n",
    "        self.attention_head_num = attention_head_num\n",
    "\n",
    "        assert embedding_dim % attention_head_num == 0, \"embedding_dim must be divisible by attention_head_num\"\n",
    "        self.head_dim = embedding_dim // attention_head_num\n",
    "\n",
    "        (\n",
    "            self.key_parameters,\n",
    "            self.output_parameters,\n",
    "            self.query_parameters,\n",
    "            self.value_parameters\n",
    "        ) = (\n",
    "            Linear(\n",
    "                config,\n",
    "                in_features=embedding_dim,\n",
    "                out_features=embedding_dim,\n",
    "                bias=False\n",
    "            )\n",
    "            for _ in range(4)\n",
    "        )\n",
    "\n",
    "        self.pre_norm = RMSNorm(config, embedding_dim)\n",
    "\n",
    "\n",
    "    def forward(self, batch: dict) -> dict:\n",
    "        features = batch[\"features\"] # (e, b, s, d)\n",
    "\n",
    "        features_normalized = self.pre_norm(features) # (e, b, s, d)\n",
    "\n",
    "        keys, queries, values = (\n",
    "            parameters(features_normalized).reshape(\n",
    "                *features.shape[:-1],\n",
    "                self.attention_head_num,\n",
    "                self.head_dim\n",
    "            ).transpose(-2, -3)\n",
    "            for parameters in (\n",
    "                self.key_parameters,\n",
    "                self.query_parameters,\n",
    "                self.value_parameters\n",
    "            )\n",
    "        ) # (e, b, a, s, k)\n",
    "\n",
    "        mask = batch[\"mask\"] # (e, b, s)\n",
    "        mask = (mask.unsqueeze(-1) & mask.unsqueeze(-2)) # (e, b, s, s)\n",
    "        mask = mask.unsqueeze(-3) # (e, b, 1, s, s)\n",
    "\n",
    "        aggregated_values = F.scaled_dot_product_attention(\n",
    "            queries,\n",
    "            keys,\n",
    "            values,\n",
    "            attn_mask=mask,\n",
    "        )\n",
    "\n",
    "        update = self.output_parameters(\n",
    "            aggregated_values.transpose(-2, -3).reshape(*features.shape)\n",
    "        ) # (e, b, s, d)\n",
    "\n",
    "        features = features + update # (e, b, s, d)\n",
    "\n",
    "        batch[\"features\"] = features\n",
    "\n",
    "        return batch\n",
    "\n",
    "embedding = Embedding(\n",
    "    config,\n",
    "    embedding_dim=config[\"embedding_dim\"],\n",
    "    vocabulary_size=len(id2token)\n",
    ")\n",
    "attention_block = AttentionBlock(\n",
    "    2,\n",
    "    config,\n",
    "    16\n",
    ")\n",
    "minibatch = embedding(minibatch)\n",
    "print(\"After embedding:\")\n",
    "for key, value in minibatch.items():\n",
    "    print(key, value.shape)\n",
    "\n",
    "minibatch = attention_block(minibatch)\n",
    "print(\"After attention block:\")\n",
    "for key, value in minibatch.items():\n",
    "    print(key, value.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cefae545",
   "metadata": {},
   "source": [
    "### Feedforward Block\n",
    "\n",
    "In the feedforward block, we just have to take an SNN via `get_mlp` and weave it through a residual connection and pre-normalization.\n",
    "\n",
    "For testing, make a block, send the minibatch through it and check keys and shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c82290f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeedforwardBlock(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Ensemble-ready transformer feedforward block. We let the hidden dimension be `4 * embedding_dim`.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        `\"device\"` : `str`\n",
    "            The device to store parameters on.\n",
    "        `\"ensemble_shape\"` : `tuple[int]`\n",
    "            The shape of the ensemble of affine transformations\n",
    "            the model represents.\n",
    "\n",
    "    Calling\n",
    "    -------\n",
    "    The expected input is a dictionary that has the following\n",
    "    key-tensor pair:\n",
    "    `\"features\"`\n",
    "        The feature tensor, of shape\n",
    "        `(..., sequence_dim, feature_dim)`\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        config: dict,\n",
    "        embedding_dim: int,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.config = config\n",
    "\n",
    "        self.mlp = get_mlp(\n",
    "            config,\n",
    "            embedding_dim,\n",
    "            embedding_dim,\n",
    "            hidden_layer_num=1,\n",
    "            hidden_layer_size=embedding_dim * 4,\n",
    "        )\n",
    "\n",
    "        self.pre_norm = RMSNorm(config, embedding_dim)\n",
    "\n",
    "\n",
    "    def forward(self, batch: dict) -> dict:\n",
    "        features = batch[\"features\"] # (e, b, s, d)\n",
    "\n",
    "        features_normalized = self.pre_norm(features) # (e, b, s, d)\n",
    "\n",
    "        update = self.mlp(features_normalized) # (e, b, s, d)\n",
    "\n",
    "        features = features + update # (e, b, s, d)\n",
    "\n",
    "        batch[\"features\"] = features\n",
    "\n",
    "        return batch\n",
    "\n",
    "feedforward_block = FeedforwardBlock(\n",
    "    config,\n",
    "    16\n",
    ")\n",
    "minibatch = feedforward_block(minibatch)\n",
    "print(\"After feedforward block:\")\n",
    "for key, value in minibatch.items():\n",
    "    print(key, value.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04abc965",
   "metadata": {},
   "source": [
    "### Output Head\n",
    "\n",
    "Note that the output of a transformer has shape end `(sequence_length, embedding_dim)`. We need a so-called output head, not to be confused with attention heads, to transform these to binary classification logits. We can first use a `MeanPool` layer and then a 16-to-1 affine transformation.\n",
    "\n",
    "For the latter, we can use `Linear` with one caveat: The layer, as implemented in Notebook 0706, takes as input and outputs a plain tensor of features. So that we can make the full models a `torch.Sequential` of the layers we created, we need to wrap it in a new layer that takes as input and outputs a dictionary with a `\"features\"` key.\n",
    "\n",
    "With this, you can create your mini-transformer as `torch.Sequential` of the embedding, the attention block, the feedforward, the mean pool, and the wrapped linear. Note that larger transformers use multiple attention block-feedforward block pairs.\n",
    "\n",
    "With the model in hand, you can again feed through the minibatch and check keys and sizes. Afterwards, you can use `evaluate_model` to get random initialization metrics, just like in Notebook 0721."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9d047a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeatureModel(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Wraps a model that takes a feature tensor and returns a feature tensor.\n",
    "    The resulting model can be used in a training loop that expects a dictionary of tensors as input and output, where the feature tensor is stored under the key `\"features\"`.\n",
    "    \"\"\"\n",
    "    def __init__(self, model: torch.nn.Module):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "\n",
    "\n",
    "    def forward(self, batch: dict) -> dict:\n",
    "        features = batch[\"features\"]\n",
    "        features = self.model(features)\n",
    "        batch[\"features\"] = features\n",
    "\n",
    "        return batch\n",
    "\n",
    "model = torch.nn.Sequential(\n",
    "    embedding,\n",
    "    attention_block,\n",
    "    feedforward_block,\n",
    "    MeanPool(config),\n",
    "    FeatureModel(Linear(config, 16, 1))\n",
    ")\n",
    "\n",
    "minibatch = model(minibatch)\n",
    "print(\"After model:\")\n",
    "for key, value in minibatch.items():\n",
    "    print(key, value.shape)\n",
    "\n",
    "evaluate_model(\n",
    "    config,\n",
    "    valid_tokenized,\n",
    "    get_binary_accuracy,\n",
    "    model,\n",
    "    \"label\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35c54cfa",
   "metadata": {},
   "source": [
    "Now you can train the model just like in Notebook 0721."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31b8dee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = AdamW(model.parameters())\n",
    "log = train_supervised(\n",
    "    config,\n",
    "    train_tokenized,\n",
    "    valid_tokenized,\n",
    "    get_binary_cross_entropy,\n",
    "    get_binary_accuracy,\n",
    "    model,\n",
    "    optimizer,\n",
    "    dataset_size_key=\"label\",\n",
    "    target_key=\"label\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3093f6b6",
   "metadata": {},
   "source": [
    "Talk about much ado about nothing. As a transformer without positional encoding is permutation equivariant, it can't really improve upon the word vector approach. Tomorrow, we'll add positional encoding."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fe50705",
   "metadata": {},
   "source": [
    "## License\n",
    "\n",
    "This work is licensed under Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International. To view a copy of this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dml",
   "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.13.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}