{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f4aa6d6",
   "metadata": {},
   "source": [
    "# Embedding Vectors\n",
    "\n",
    "## Setup\n",
    "\n",
    "### Install\n",
    "\n",
    "Since today, we'll tokenize texts to words, for simplicity, we will use the Moses tokenizer https://github.com/luismsgomes/mosestokenizer.\n",
    "\n",
    "You can install it via the following line, to be used with your `mamba` environment activated:\n",
    "```bash\n",
    "pip install mosestokenizer\n",
    "```\n",
    "\n",
    "If you run into an error with this tokenizer, you can do\n",
    "```bash\n",
    "pip install sacremoses\n",
    "```\n",
    "instead."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "998f1a26",
   "metadata": {},
   "source": [
    "### Imports\n",
    "\n",
    "Import `defaultdict`, `Callable`, `Iterable`, `datasets`, `matplotlib.pyplot` as `plt`, `math`, `os`, `torch`, `torch.nn.functional` as `F`, `tqdm`, and `Optional`.\n",
    "\n",
    "2. Also import the new module `mosestokenizer`, or if you ran into an error with that, `sacremoses`.\n",
    "\n",
    "Moreover, import the following:\n",
    "1. The function `get_random_reshuffler`, that you wrote in Notebook 0623,\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 functions `pbt_init` and `pbt_update`, that you wrote in Notebook 0707, and\n",
    "1. The classes `AdamW` and `Optimizer`, that you wrote in Notebook 0708."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98a5f96",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f08a5ff0",
   "metadata": {},
   "source": [
    "### Configuration\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",
    "    Make this `(16,)`.\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",
    "    Make this a `64`.\n",
    "- `\"minibatch_size_eval\"`: `int`  \n",
    "    On my home computer, I can make this `128`.\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": "f449f42e",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d71a33b9",
   "metadata": {},
   "source": [
    "Set the `torch` pseudo-random number generation seed, as per the configuration dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b25f9805",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7857b425",
   "metadata": {},
   "source": [
    "## Data Preprocessing\n",
    "\n",
    "### Load and Split Dataset\n",
    "\n",
    "Load the dataset and print it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6049c77b",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7f707da",
   "metadata": {},
   "source": [
    "Looks like we have a test split as big as the train split. Let's split the test split then to a valid-test split, using `datasets.Dataset.train_test_split`. Don't forget to set the seed using `get_seed`.\n",
    "\n",
    "Assign the original train and the new validation split to variables. Print the splits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e875713",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86a191a1",
   "metadata": {},
   "source": [
    "Print the `\"text\"` field of the first 3 entries of the train set. These are the sort of texts we'll work with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cae2473f",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "feeb05ff",
   "metadata": {},
   "source": [
    "### Storing Sequential Data\n",
    "\n",
    "We will preprocess the data by transforming words to token indices. We will have to store sequences of varying lengths. To save space, we will store sequential data in 2 vectors:\n",
    "1. `token_ids`: this stores token IDs of all documents in order.\n",
    "2. `indptr`: the token IDs of the `i`-th document are `token_ids[indptr[i]:indptr[i+1]]`.\n",
    "\n",
    "Note: The variable name `indptr` stands for *index pointer*. It comes from the naming convention of `scipy` sparse matrices in Compressed Sparse Row format: https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html\n",
    "\n",
    "1. Write the function below.\n",
    "2. Get its output on the validation set with no optional arguments.\n",
    "3. Print the sizes of the tensors and the vocabulary.\n",
    "3. Using the 2 vectors in the output, get the token IDs of the 100-th validation entry.\n",
    "4. Using the ID-to-token list of the output, convert the token IDs you got in part 3. to a document. Print it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfa75952",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_corpus(\n",
    "    config: dict,\n",
    "    corpus: list[str],\n",
    "    fixed_vocabulary: Optional[bool] = None,\n",
    "    id2token: Optional[list] = None,\n",
    "    token2id: Optional[dict] = None\n",
    ") -> tuple[\n",
    "    torch.Tensor,\n",
    "    torch.Tensor,\n",
    "    list,\n",
    "    dict\n",
    "]:\n",
    "    \"\"\"\n",
    "    Tokenize a corpus using `mosestokenizer`.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pair:\n",
    "        `\"device\"` : `int | str | torch.device`\n",
    "            Device to store tensors on.\n",
    "    corpus : `list[str]`\n",
    "        The document corpus.\n",
    "    fixed_vocabulary : `bool`, optional\n",
    "        Determines if the vocabulary given in the\n",
    "        `id2token` and `token2id` keyword arguments can be augmented\n",
    "        or tokens not in the vocabulary should be discarded.\n",
    "        If not given, then it will be `id2token is not None`.\n",
    "    id2token : `list[str]`\n",
    "        The mapping from token IDs to tokens.\n",
    "        If `not fixed_vocabulary`, then it is augmented\n",
    "        as new tokens are extracted from the corpus. Default: `[]`.\n",
    "    token2id : `dict`\n",
    "        The mapping from tokens to token IDs.\n",
    "        If `not fixed_vocabulary`, then it is augmented\n",
    "        as new tokens are extracted from the corpus. Default: `{}`.\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    The quadruple of:\n",
    "    1. The vector of token IDs,\n",
    "    2. the vector of document index pointers,\n",
    "    3. the list `id2token` and\n",
    "    4. the dictionary `token2id`.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bedbb18a",
   "metadata": {},
   "source": [
    "Via another call to `tokenize_corpus`, get the training token IDs and training index pointers, while augmenting the vocabulary with new tokens from the training corpus.\n",
    "\n",
    "Once again, print sizes and the text of the 100-th entry as extracted from the tokenized data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fe07ba1",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ccf1d03",
   "metadata": {},
   "source": [
    "### Create Dataset Dictionaries\n",
    "\n",
    "For more versatility, we'll make preprocessed datasets, and minibatches thereof, dictionaries. So, make training and validation datasets using the token ids and index pointers you extracted and the labels in the downloaded dataset. The latter should be converted to tensors.\n",
    "\n",
    "Print keys and value shapes of the dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "832f8e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afe01e86",
   "metadata": {},
   "source": [
    "Save the preprocessed splits into a file. I recommend saving the dictionary with `\"train\"`, `\"valid\"`, `id2token`, and `token2id` fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f2f2e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a1a6f6",
   "metadata": {},
   "source": [
    "## Dataloaders on Sequential Data\n",
    "\n",
    "Now that we have dataset entries of varying length stored via `token_ids` and `indptr`, we need to update the `get_minibatch` function.\n",
    "\n",
    "### `is_ensembled` and `to_ensembled`\n",
    "\n",
    "First of all, write the following functions, that respectively check if a tensor has ensemble dimensions and converts a tensor to have them if not.\n",
    "\n",
    "Test the functions on the validation label tensor:\n",
    "1. `is_ensembled` should be false for it.\n",
    "2. Applying `to_ensembled`, it should get ensemble dimensions. You can check this via printing the shape of the original and the transformed tensor.\n",
    "3. `is_ensembled` should be true for the transformed tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e4b802a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_ensembled(\n",
    "    ensemble_shape: tuple[int],\n",
    "    tensor: torch.Tensor\n",
    ") -> bool:\n",
    "    \"\"\"\n",
    "    We view `tensor` as *ensembled* if it is prefixed by `ensemble_shape`,\n",
    "    that is its slice of the first `len(ensemble_shape)` entries\n",
    "    is `ensemble_shape`.\n",
    "\n",
    "    This function checks this condition.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "def to_ensembled(\n",
    "    ensemble_shape: tuple[int],\n",
    "    tensor: torch.Tensor\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    We say that a tensor is *ensembled*,\n",
    "    if its shape starts by the ensemble shape.\n",
    "\n",
    "    This function converts a tensor to an ensembled tensor.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9a0e81b",
   "metadata": {},
   "source": [
    "### `get_minibatch`\n",
    "\n",
    "Time to write a function that, given a dataset, minibatch IDs, and optionally an ensemble shape, outputs an appropriate minibatch dictionary. There will be 3 type of fields in there:\n",
    "1. For fields without a sequential dimension, that is that are *arraylike* (my terminology), a tensor where we index into the entry dimension with the minibatch IDs, as before.\n",
    "2. For fields with a sequential dimension, that is that are *sequential* a tensor with sequential dimension the maximum of the lengths of the entries at the minibatch IDs\n",
    "3. `mask`: A boolean tensor that says how much of the padded sequential dimensions are actually taken up by each entry.\n",
    "\n",
    "To fix definition: We'll call *dataset size* the number of entries in a dataset, and *total number of sequence entries* the sum of the lengths of the sequences in a dataset. Print both for both the training and validation sets, using the `indptr` fields. While we're at it, print also the vocabulary size, that is the number of unique words we've collected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e06c83f",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c07c4119",
   "metadata": {},
   "source": [
    "1. Note that the minibatch IDs we receive may (training dataloader) or may not (evaluation) have ensemble dimensions. Therefore, first of all, if we received an ensemble shape, we should make sure that it is ensembled by conversion, if needed.\n",
    "\n",
    "2. Now check if the dataset has an `indptr` field. If not, you can just index into each tensor as before.\n",
    "\n",
    "3. From now on, let's assume the dataset does have an `indptr` field.\n",
    "    1. Let's flatten the minibatch IDs. It will be easier to work with them; we'll unflatten tensor upon adding them to the minibatch dictionary.\n",
    "    1. Using indptr and minibatch IDs, you can get the lengths of the dataset entries at the minibatch IDs. In particular, we can select the maximum length.\n",
    "    1. Thus, we can make the sequential tensors in the minibatch have shape `ensemble_shape + (minibatch_size, length_max) + feature_shape`.\n",
    "    1. We can populate these tensors by indexing into the respective tensors in the dataset by an index tensor of shape `ensemble_shape + (minibatch_size, length_max)` that at given `e + (i,) in ensemble_shape + (minibatch_size,)` is `indptr[e + (i,)], indptr[e + (i,)] + 1, ..., indptr[e + (i,)] + length_max - 1`, clamped at total number of sequence entries - 1. You can get this index tensor via broadcasted addition.\n",
    "    1. We also need a mask tensor of shape `ensemble_shape + (minibatch_size, length_max)` that say which entries are not padding entries. For this, you can use the entrywise length values.\n",
    "    1. Now you are ready to create the minibatch dictionary. Don't forget to unflatten tensors.\n",
    "\n",
    "To test `get_minibatch`, get a minibatch index tensor via `get_random_reshuffler`. Print keys and value shapes of the minibatch dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7b6ac94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_minibatch(\n",
    "    dataset: dict,\n",
    "    minibatch_ids: torch.Tensor,\n",
    "    ensemble_shape: Optional[tuple[int]] = None\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Get a minibatch from the dataset. The minibatch IDs\n",
    "    can be ensembled. If the dataset has an `indptr` field,\n",
    "    then sequential data is also minibatched properly.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    dataset : `dict`\n",
    "        The dataset to get the minibatch from.\n",
    "    minibatch_ids : `torch.Tensor`\n",
    "        The IDs of the entries to include in the minibatch.\n",
    "        Can be of shape `ensemble_shape + (minibatch_size,)`\n",
    "        or just `(minibatch_size,)`.\n",
    "    ensemble_shape : `tuple[int]`, optional\n",
    "        The ensemble shape. If given, then the minibatch IDs\n",
    "        are converted to ensembled minibatch IDs. Default: `None`.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    A dictionary representing the minibatch, with the following key-value pairs:\n",
    "    For each key in the dataset, a tensor of the corresponding entries\n",
    "    at the minibatch IDs. If the dataset has an `indptr` field,\n",
    "    then for sequential fields, the returned tensor has shape\n",
    "    `ensemble_shape + (minibatch_size, length_max) + feature_shape`,\n",
    "    where `length_max` is the maximum length of the sequences\n",
    "    at the minibatch IDs, and `feature_shape` is the shape\n",
    "    of the feature dimensions of the corresponding dataset tensor,\n",
    "    excluding the sequential dimension.\n",
    "    `\"mask\"` : `torch.Tensor`\n",
    "        A boolean tensor of shape\n",
    "        `ensemble_shape + (minibatch_size, length_max)`,\n",
    "        that signals which entries are not padding elements.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce8f2a95",
   "metadata": {},
   "source": [
    "Now that you wrote a new `get_minibatch` function, you can update the `get_dataloader_random_reshuffle` function that you wrote in Notebook 0623 to get a dataloader that supports sequential data.\n",
    "\n",
    "Get a train dataloader, get a minibatch, print keys and value shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d203ae3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataloader_random_reshuffle(\n",
    "    config: dict,\n",
    "    dataset: dict,\n",
    "    dataset_size_key=\"label\"\n",
    ") -> Iterable[dict]:\n",
    "    \"\"\"\n",
    "    Given a dataset with optionally sequential data, this function\n",
    "    creates a random reshuffling (without replacement) dataloader\n",
    "    that yields minibatch dictionaries indefinitely.\n",
    "    Support arbitrary ensemble shapes.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : dict\n",
    "        Configuration dictionary. Required keys:\n",
    "        \"device\" : str\n",
    "            The device to store tensors on.\n",
    "        ensemble_shape : tuple[int]\n",
    "            The required ensemble shapes of the outputs.\n",
    "        minibatch_size : int\n",
    "            The size of the minibatches.\n",
    "    dataset : dict\n",
    "        The dataset to get the minibatches from. It can contain sequential data, in concatenated form, in which case it should have an `indptr` field.\n",
    "    dataset_size_key : str, optional\n",
    "        The key in the dataset that maps to a tensor with first dimension\n",
    "        the dataset size. Default: `\"label\"`    \n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    An iterable that yields minibatch dictionaries indefinitely.\n",
    "    Each minibatch dictionary has the same keys as the dataset,\n",
    "    with values corresponding to the entries at the minibatch IDs,\n",
    "    except that for sequential data, the returned tensors have shape\n",
    "    `ensemble_shape + (minibatch_size, length_max) + feature_shape`,\n",
    "    where `length_max` is the maximum length of the sequences\n",
    "    at the minibatch IDs, and `feature_shape` is the shape of the\n",
    "    feature dimensions of the corresponding dataset tensor,\n",
    "    excluding the sequential dimension.\n",
    "    In addition, if the dataset has sequential data,\n",
    "    then each minibatch dictionary has a `\"mask\"` key,\n",
    "    with a boolean tensor value of shape\n",
    "    `ensemble_shape + (minibatch_size, length_max)`,\n",
    "    that signals which entries are not padding elements.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d63597a",
   "metadata": {},
   "source": [
    "## Layers on Minibatch Dictionaries\n",
    "\n",
    "From now on, we'll use layers that expect and output minibatch dictionaries. Today, as discussed in lecture, we'll use 1-dimensional token embedding vectors that can be transformed to binary classification logits after a mean pool."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2156e98",
   "metadata": {},
   "source": [
    "### `Embedding`\n",
    "\n",
    "First of all, let's create a token embedding layer. Its parameters will be a matrix of one row vector of `embedding_dim` per vocabulary entry. Then upon reception of a minibatch dictionary, it will expect a `token_ids` field and it will output an updated dictionary with a new `features` field that has shape `token_ids.shape + (embedding_dim,)`, where each token ID is replaced by the corresponding token embedding vector.\n",
    "\n",
    "Write the layer, instantiate it with `embedding_dim=1`, feed it the minibatch, and write keys and value shapes of the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b10d9e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embedding(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Ensemble-ready embedding.\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",
    "    embedding_dim : `int`\n",
    "        The number of embedding dimensions.\n",
    "    vocabulary_size : `int`\n",
    "        The number of vocabulary entries.\n",
    "\n",
    "    Calling\n",
    "    -------\n",
    "    Instance calls require one positional argument:\n",
    "    batch : `dict`\n",
    "        The input data dictionary. Required key:\n",
    "        `\"token_ids\"` : `torch.Tensor`\n",
    "            Tensor of token IDs. It is required to be one of the following shapes:\n",
    "            1. `ensemble_shape + batch_shape`\n",
    "            2. `batch_shape`\n",
    "\n",
    "        Upon a call, the model thinks we're in the first case\n",
    "        if the first `len(ensemble_shape)` many entries of the\n",
    "        shape of the input tensor is `ensemble_shape`.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0cfd744",
   "metadata": {},
   "source": [
    "### `MeanPool`\n",
    "\n",
    "Let's now make a mean pool layer. The novelty with respect to the pooling operations for CNNs is that now we are pooling over various sequence lenghts. To this end, you can make a broadcasted product of the feature tensor with the mask tensor, and divide by the sum of the mask tensor along the sequence dimension. The booleans will get cast to floats and the sum will be understood as the sequence length.\n",
    "\n",
    "Write the `MeanPool` class below, instantiate it, and apply it to the minibatch output by the embedding layer instantiated above. Loop over the entries in the minibatch it outputs and print keys and value shapes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "534e789e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MeanPool(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    This module performs mean pool of sequential data\n",
    "    along the sequential dimension.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pair:\n",
    "        `\"ensemble_shape\"` : `tuple[int]`\n",
    "            The shape of the ensemble of models\n",
    "            that process the data.\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",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8d31ed7",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Now you just need to make superficial updates to `evaluate_model` and `train_supervised`, respectively written in Notebooks 0707 and 0720, to reflect that now our datasets are dictionaries and our models expect and output dictionaries.\n",
    "\n",
    "Once you're done updating the functions, you can traid 1D token embedding vectors on IMDB. Assign the log dictionary to a variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1869d323",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(\n",
    "    config: dict,\n",
    "    dataset: dict,\n",
    "    get_metric: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],\n",
    "    model: torch.nn.Module,\n",
    "    target_key=\"target\"\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Evaluate a model on a supervised dataset.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pair:\n",
    "        `\"minibatch_size_eval\"` : `int`\n",
    "            Size of consecutive minibatches to take from the dataset.\n",
    "            To be set according to RAM or GPU memory capacity.\n",
    "    dataset : `dict`\n",
    "        The dataset to evaluate the model on.\n",
    "    get_metric : `Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`\n",
    "        Function to get the metric from a pair of\n",
    "        predicted and target value tensors.\n",
    "    model : `torch.nn.Module`\n",
    "        The model to evaluate.\n",
    "    target_key : `str`, optional\n",
    "        The key mapped to the target value tensor in the dataset.\n",
    "        Default: `\"target\"`\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cbef706",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_supervised(\n",
    "    config: dict,\n",
    "    dataset_train: dict,\n",
    "    dataset_valid: dict,\n",
    "    get_loss: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],\n",
    "    get_metric: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],\n",
    "    model: torch.nn.Module,\n",
    "    optimizer: Optimizer,\n",
    "    dataset_size_key=\"label\",\n",
    "    target_key=\"target\",\n",
    ") -> dict:\n",
    "    \"\"\"\n",
    "    Population-based training on a supervised learning task.\n",
    "    Tuned hyperparameters are given by raw values and transformations.\n",
    "    This way, the hyperparameters are perturbed by\n",
    "    additive noise on raw values.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    config : `dict`\n",
    "        Configuration dictionary. Required key-value pairs:\n",
    "        `\"ensemble_shape\"` : tuple[int]\n",
    "            Ensemble shape. We assume this is a 1-dimensional tuple\n",
    "            with dimensions the population size.\n",
    "        `\"hyperparameter_raw_init_distributions\"` : `dict`\n",
    "            Dictionary that maps tuned hyperparameter names\n",
    "            to `torch.distributions.Distribution` of raw hyperparameter values.\n",
    "            Required keys:\n",
    "            `\"learning_rate\"`:\n",
    "                The learning rate of stochastic gradient descent.\n",
    "        `\"hyperparameter_raw_perturbs\"` : `dict`\n",
    "            Dictionary that maps tuned hyperparameter names\n",
    "            to `torch.distributions.Distribution` of additive noise.\n",
    "        `\"hyperparameter_transforms\"` : `dict`\n",
    "            Dictionary that maps tuned hyperparameter names\n",
    "            to transformations of raw hyperparameter values.\n",
    "        `\"improvement_threshold\"` : `float`\n",
    "            A new metric score has to be this much better\n",
    "            than the previous best to count as an improvement.\n",
    "        `\"minibatch_size\"` : `int`\n",
    "            Minibatch size to use in a training step.\n",
    "        `\"minibatch_size_eval\"` : `int`\n",
    "            Minibatch size to use in evaluation.\n",
    "            On CPU, should be about the same as `minibatch_size`.\n",
    "            On GPU, should be as big as possible without\n",
    "            incurring an Out of Memory error.\n",
    "        `\"pbt\"` : `bool`\n",
    "            Whether to use PBT updates in validations.\n",
    "            If `False`, the algorithm just samples hyperparameters at start,\n",
    "            then keeps them constant.\n",
    "        `\"steps_num\"` : `int`\n",
    "            Maximum number of training steps.\n",
    "        `\"steps_without_improvement`\" : `int`\n",
    "            If the number of training steps without improvement\n",
    "            exceeds this value, then training is stopped.\n",
    "        `\"valid_interval\"` : `int`\n",
    "            Frequency of evaluations, measured in number of training steps.\n",
    "        `\"welch_confidence_level\"` : `float`\n",
    "            The confidence level in Welch's t-test\n",
    "            that is used in determining if a population member\n",
    "            is to be replaced by another member with perturbed hyperparameters.\n",
    "        `\"welch_sample_size\"` : `int`\n",
    "            The last this many validation metrics are used\n",
    "            in Welch's t-test.\n",
    "    `dataset_train` : `dict`\n",
    "        The training dataset.\n",
    "    `dataset_valid` : `dict`\n",
    "        The validation dataset.\n",
    "    `get_loss` : `Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`\n",
    "        A function that maps a pair of predicted and target value tensors\n",
    "        to a tensor of losses per ensemble member.\n",
    "    `get_metric` : `Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`\n",
    "        A function that maps a pair of predicted and target value tensors\n",
    "        to a tensor of metrics per ensemble member.\n",
    "        We assume a greater metric is better.\n",
    "    `model` : `torch.nn.Module`\n",
    "        The model ensemble to tune.\n",
    "    `optimizer` : `Optimizer`\n",
    "        An optimizer that tracks the parameters of `model`.\n",
    "    `dataset_size_key` : `str`, optional\n",
    "        The key in the datasets that maps to a tensor\n",
    "        whose first dimension is the dataset size.\n",
    "        Default: `\"label\"`\n",
    "    `target_key` : `str`, optional\n",
    "        The key mapped to the target value tensor in the datasets.\n",
    "        Default: `\"target\"`\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    An output dictionary with the following key-value pairs:\n",
    "        `\"best parameters\"` : `dict`  \n",
    "            The state dictionary of the model with the best metric\n",
    "            encountered during training.\n",
    "        `\"source mask\"` : `torch.Tensor`\n",
    "            The source masks of population members\n",
    "            that were replace by other members in a PBT update\n",
    "        `\"target indices\"` : `torch.Tensor`\n",
    "            The indices of population members\n",
    "            that the member where the source mask is to were replaced with.\n",
    "        `\"validation metric\"` : `torch.Tensor`\n",
    "            The validation metrics at evaluation steps.\n",
    "\n",
    "        In addition, for each tuned hyperparameter name,\n",
    "        we include a `torch.Tensor` of values per update.\n",
    "    \"\"\"\n",
    "    raise NotImplementedError\n",
    "\n",
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e50641a3",
   "metadata": {},
   "source": [
    "Note that `log[\"best_parameters\"]` has the parameters of the best model seen during training. In particular, the `\"embedding\"` field of that is the token embedding matrix of the best model. Using the vocabulary, list the words with the highest 10 and lowest 10 token embedding values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d34c2bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca571ec9",
   "metadata": {},
   "source": [
    "## Dataset References\n",
    "\n",
    "[1] Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts: *Learning Word Vectors for Sentiment Analysis*, 2011. Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies (ACL 2011), pages 142--150. [link](https://aclanthology.org/P11-1015/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aae9f3d",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
