token_and_embed.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c44d08d1",
   "metadata": {},
   "source": [
    "```\n",
    "- Classifier       - takes in text OR image                           outputs a number (category).\n",
    "                        newsgroups OR MNIST\n",
    "- Regressor        - takes in numbers                                 outputs numbers.\n",
    "                        house price\n",
    "- AutoEncoder      - takes in images                                  outputs images.\n",
    "                        MNIST\n",
    "- LLM/Transformer  - takes in text                                    outputs text.\n",
    "                        Llama\n",
    "- ???              - takes in text AND numbers AND images AND audio   outputs text AND numbers AND images AND audio.\n",
    "                        Gato\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "132dbd4f",
   "metadata": {},
   "source": [
    "How would you do that?\n",
    "\n",
    "We've used CountVectorizer and TfidfVectorizer to turn text into vector embeddings.\n",
    "\n",
    "We've used ResNet to turn images into vector embeddings.\n",
    "\n",
    "We haven't used anything to combine text and images into a single embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a3dbf0a-ec67-4aab-9e44-d0fe41937a38",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import random\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tiktoken\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import tensor\n",
    "import transformers\n",
    "from transformers import GPT2Config\n",
    "from functools import partial\n",
    "\n",
    "from mugato.data import ok_vqa, four_rooms\n",
    "from mugato.tokenizer import Tokenizer\n",
    "from mugato.utils import image_transform, as_tensor\n",
    "\n",
    "random.seed(42)\n",
    "\n",
    "def to_numpy(x):\n",
    "    if isinstance(x, torch.Tensor):\n",
    "        return x.clone().detach().numpy()\n",
    "    else:\n",
    "        return x\n",
    "\n",
    "def plot_embedding(xs, ys, zs, max_x=10, max_y=10, max_z=10, colors=None):\n",
    "    if not isinstance(xs, list):\n",
    "        xs, ys, zs = [xs], [ys], [zs]\n",
    "    # Create a new figure and add a 3D subplot\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "    if colors is None:\n",
    "        colors = ['red'] * len(xs)\n",
    "    \n",
    "    # Plot the points\n",
    "    with torch.no_grad():\n",
    "        ax.scatter(to_numpy(xs), to_numpy(ys), to_numpy(zs), color=colors, marker='o')\n",
    "    \n",
    "    # Set axes limits\n",
    "    ax.set_xlim(-7, 7)\n",
    "    ax.set_ylim(-7, 7)\n",
    "    ax.set_zlim(-7, 7)\n",
    "    \n",
    "    # Label the axes\n",
    "    ax.set_xlabel('X (\"A\"-ness)')\n",
    "    ax.set_ylabel('Y (\"vowely\"-ness)')\n",
    "    ax.set_zlabel('Z (\"lettery\"-ness)')\n",
    "    \n",
    "    # Display the plot\n",
    "    plt.show()\n",
    "\n",
    "def imshow(image):\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.imshow(image.permute(1, 2, 0))\n",
    "    plt.show()\n",
    "\n",
    "vqa_dataset = ok_vqa.initialize()\n",
    "sample_index = 10\n",
    "sample = vqa_dataset['train'][sample_index]\n",
    "image = as_tensor(sample['image']) / 255\n",
    "\n",
    "four_rooms_dataset = four_rooms.initialize()\n",
    "\n",
    "tiktoken_tokenizer = tiktoken.get_encoding(\"r50k_base\")\n",
    "tokenizer = Tokenizer(tiktoken_tokenizer)\n",
    "vqa_dataset = ok_vqa.initialize()\n",
    "vqa_dataloader = ok_vqa.create_dataloader(tokenizer, batch_size=4)\n",
    "dlit = iter(vqa_dataloader);\n",
    "next(dlit); next(dlit);\n",
    "batch = next(dlit)\n",
    "\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91bf1eac-012c-4a3f-834a-8c44d442e230",
   "metadata": {},
   "source": [
    "# Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d590813-6596-4143-b9cb-78a22a94f1e7",
   "metadata": {},
   "source": [
    "Are these two phrases similar?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33a126cf-3f8c-4989-8702-1ff4ee10fe68",
   "metadata": {},
   "outputs": [],
   "source": [
    "word_1 = \"big blue ocean\"\n",
    "word_2 = \"little green frog\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d0aad1-f84f-47bb-913d-d87d07c68200",
   "metadata": {},
   "source": [
    "<details>\n",
    "\n",
    "It depends on what conceptual idea you compare them.\n",
    "\n",
    "They are both things you might see in a painting of a landscape.\n",
    "\n",
    "One you can live in, the other you can't.\n",
    "\n",
    "They both have the same number of letters\n",
    "\n",
    "One has an 'h', the other doesn't.\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a917b0ee-163a-469d-a96c-88000c79c739",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper functions to get \"features\" of our phrases.\n",
    "def count_vowels(text):\n",
    "    return len([c for c in text if c in 'aeiouy']) / 2\n",
    "\n",
    "def count_letter(letter, text):\n",
    "    return len([c for c in text if c == letter]) / 3\n",
    "\n",
    "count_L = partial(count_letter, 'l')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d99c962-5c97-4eb5-9fe6-b10f9c787c06",
   "metadata": {},
   "outputs": [],
   "source": [
    "def example_embedding_1(text):\n",
    "    \"\"\"\n",
    "    Embeds the text in a 3-dimensional space.\n",
    "    \n",
    "    [            x,                y,              z]\n",
    "\n",
    "    [number of 'l', number of vowels, length of text]\n",
    "\n",
    "    \"\"\"\n",
    "    x = count_L(text)      # number of \"l\"s\n",
    "    y = count_vowels(text) # number of vowels\n",
    "    z = len(text) / 3      # number of letters / 3\n",
    "    return x, y, z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94505b08-446f-4aea-9412-60703a0612c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_embedding_1(\"big blue ocean\"), example_embedding_1(\"little green frog\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7675b95",
   "metadata": {},
   "source": [
    "Because our embedding space is 3-dimensions, we can visualize where these phrases land in embedding space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab7c564-4163-4143-9c2e-98a7f7357412",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*example_embedding_1('big blue ocean'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4273fdff-e3fc-44e0-8615-3198bdaf391a",
   "metadata": {},
   "source": [
    "This next cell, I just want to show you what it looks like to plot multiple things in the same 3d space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c1de4e1-3f13-4507-9b88-fcf447289044",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    example_embedding_1('big blue ocean'),\n",
    "    example_embedding_1('little blue bird'),\n",
    "    example_embedding_1('big green hulk'),\n",
    "    example_embedding_1('frog'),\n",
    "), colors=[\n",
    "    \"blue\",\n",
    "    \"blue\",\n",
    "    \"green\",\n",
    "    \"green\",\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c3a9be",
   "metadata": {},
   "source": [
    "We've been using this embedding function, `example_embedding_1`, which takes a text, and returns an embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ce90970",
   "metadata": {},
   "outputs": [],
   "source": [
    "def example_embedding_1(text):\n",
    "    x = count_L(text)      # number of \"l\"s\n",
    "    y = count_vowels(text) # number of vowels\n",
    "    z = len(text) / 3      # number of letters / 3\n",
    "    return x, y, z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15ee8eb0",
   "metadata": {},
   "source": [
    "\n",
    "How do they really do it, in the real world?\n",
    "\n",
    "- CountVectorizer\n",
    "- TfidfVectorizer\n",
    "- Word2Vec\n",
    "- ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f50a849",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "1b0e39db",
   "metadata": {},
   "source": [
    "## \"Embedding Layers\" (nn.Embedding)\n",
    "\n",
    "An ***\"embedding layer\"*** is a thing that ***takes something*** (a token, a document, etc...) ***and returns an a list of numbers*** (an embedding)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c12034",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = transformers.AutoModel.from_config(GPT2Config())\n",
    "model  # First layer: Embedding(50257, 768)  <- What is the 50257? What is the 768?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23692720",
   "metadata": {},
   "source": [
    "Move this down maybe?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8613455b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = transformers.AutoModel.from_config(GPT2Config(vocab_size=50257, n_embd=3, n_head=1))\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65caf401",
   "metadata": {},
   "source": [
    "We can make our own embedding layer. (Remember the Gato paper, we want to take text, images, actions, button presses, etc... and turn them into embeddings.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fac1adfe",
   "metadata": {},
   "source": [
    "- Again, what is an \"embedding layer\"?\n",
    "- How would you make a text embedding layer?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f82399eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocabulary = [\n",
    "    'big',\n",
    "    'little',\n",
    "    'blue',\n",
    "    'green',\n",
    "    'ocean',\n",
    "    'bird',\n",
    "    'frog',\n",
    "    'hulk',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9b1e96d-bd9e-442d-9505-c22aa0a7d34f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a dictionary where each word (token) is the key, and each value is 3 random numbers.\n",
    "embedding_layer = {}\n",
    "for i, word in enumerate(vocabulary):\n",
    "    embedding_layer[word] = torch.randn(3)  # Example: embedding_layer['big'] = [0.12, 1.5, -0.3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5a30152",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_layer['big']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecc1c61e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiktoken_tokenizer = tiktoken.get_encoding(\"r50k_base\")\n",
    "tiktoken_tokenizer.encode('big blue ocean')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "295b269d",
   "metadata": {},
   "source": [
    "A new plot, like the one above, but instead of defining the meaning/concept of each dimension (numer of l's, number of vowels, length of text), we are initializing the values of each dimension randomly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c27e3eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    embedding_layer['big'],\n",
    "    embedding_layer['green'],\n",
    "    embedding_layer['hulk'],\n",
    "    embedding_layer['little'],\n",
    "    embedding_layer['blue'],\n",
    "    embedding_layer['bird'],\n",
    "), colors=[\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74cd88e9-65b6-475a-a915-95f69b3023b6",
   "metadata": {},
   "source": [
    "A brief aside:\n",
    "\n",
    "We can \"combine\" embeddings to create new embeddings.\n",
    "\n",
    "- How would you combine the embeddings of \"little\" and \"blue\"?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d50f577b",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "519e1994",
   "metadata": {},
   "outputs": [],
   "source": [
    "def combine_embeddings(embedding_1, embedding_2):\n",
    "    return (embedding_1 + embedding_2) / 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1959b011",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    embedding_layer['blue'],\n",
    "    embedding_layer['green'],    \n",
    "    combine_embeddings(embedding_layer['blue'], embedding_layer['green']),\n",
    "), colors=[\n",
    "    'blue',\n",
    "    'green',\n",
    "    'purple',\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd1f612d",
   "metadata": {},
   "source": [
    "Every new state of the art machine learning algorithm (Transformer, ResNet, etc...) is basically just a new clever way to combine embeddings.\n",
    "\n",
    "But even something as simple as taking the average works. It just doesn't work well."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f3f8175",
   "metadata": {},
   "source": [
    "- In typical English text, do you expect the word \"little\" to be near the word \"frog\"?\n",
    "\n",
    "Do you expect those two words to be close to each other in the embedding space?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b55f305",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "2aad4705-d35f-4b1e-800e-84a1072c13d0",
   "metadata": {},
   "source": [
    "# Images\n",
    "\n",
    "What about images?\n",
    "\n",
    "How would you embed an image into the same 3-dimensional embedding space that we've been using above?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e5cf7ab-d2b0-4660-a451-ee4d3d2ce374",
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow(image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b53e3bce",
   "metadata": {},
   "source": [
    "- Is an image a token?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55eda684",
   "metadata": {},
   "source": [
    "- What's an \"Embedding Layer\"?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236379ce",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "974605f8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f421c5-971f-4543-bfe5-a64d89087dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Is this an \"embedding layer\"?\n",
    "def example_image_embedding(image):\n",
    "    x = image[0].mean() * 5  # Red channel\n",
    "    y = image[1].mean() * 5  # Green channel\n",
    "    z = image[2].mean() * 5  # Blue channel\n",
    "    return torch.tensor([x, y, z])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9489ef5d",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e5fa576-ee5d-4b6a-93e1-c80e130c88e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_image_embedding(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45b3542b-4789-40cc-9e87-3a3e30e38b42",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    example_embedding_1('blue'),\n",
    "    example_embedding_1('ocean'),\n",
    "    example_image_embedding(image),\n",
    "), colors=[\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'red',\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aacc1b92",
   "metadata": {},
   "source": [
    "# \"Discrete\" button presses"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07ae06aa",
   "metadata": {},
   "source": [
    "How would you embed an Atari button press (up, down, left, right)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95db07ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "FORWARD = tensor(0)\n",
    "BACKWARD = tensor(1)\n",
    "LEFT = tensor(2)\n",
    "RIGHT = tensor(3)\n",
    "FORWARD, BACKWARD, LEFT, RIGHT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ee5868b-969e-469e-9aee-409819d8cd69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Effectively the same thing as:\n",
    "# atari_embedding = {\n",
    "#     FORWARD:    torch.randn(3),\n",
    "#     BACKWARD:  torch.randn(3),\n",
    "#     LEFT:  torch.randn(3),\n",
    "#     RIGHT: torch.randn(3),\n",
    "# }\n",
    "atari_embedding = nn.Embedding(4, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cccee4bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "atari_embedding(FORWARD)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf91a154",
   "metadata": {},
   "source": [
    "Ok.\n",
    "\n",
    "Now we can embed text, images, and discrete values like button presses.\n",
    "\n",
    "How do we embed continuous values?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3df65cf7",
   "metadata": {},
   "source": [
    "# Continuous robotic arm movements\n",
    "\n",
    "How would you encode the angular position of a robotic arm?\n",
    "\n",
    "The position could be any floating point number between 0 and 360.\n",
    "\n",
    "`nn.Embedding(..., 3)`?\n",
    "\n",
    "There's an infinite number of possible positions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecb76d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize_position(position):\n",
    "    if position >= 0 and position < 90:\n",
    "        return tensor(0)\n",
    "    elif position >= 90 and position < 180:\n",
    "        return tensor(1)\n",
    "    elif position >= 180 and position < 270:\n",
    "        return tensor(2)\n",
    "    elif position >= 270 and position <= 360:\n",
    "        return tensor(3)\n",
    "\n",
    "continuous_embedding = nn.Embedding(4, 3)\n",
    "plot_embedding(*zip(\n",
    "    example_embedding_1('blue'),\n",
    "    example_image_embedding(image),\n",
    "    atari_embedding(FORWARD),\n",
    "    continuous_embedding(tokenize_position(45.8992)),\n",
    "), colors=[\n",
    "    'red',\n",
    "    'green',\n",
    "    'blue',\n",
    "    'black',\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "920e312a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "9135c5f0",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec321c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "BLUE_IMAGE = torch.ones(3, 224, 224)\n",
    "BLUE_IMAGE[0, :, :] *= 0.25\n",
    "BLUE_IMAGE[1, :, :] *= 0.25\n",
    "BLUE_IMAGE[2, :, :] *= 0.75\n",
    "BLUE_IMAGE = nn.Parameter(BLUE_IMAGE, requires_grad=True)\n",
    "GREEN_IMAGE = torch.ones(3, 224, 224)\n",
    "GREEN_IMAGE[0, :, :] *= 0.25\n",
    "GREEN_IMAGE[1, :, :] *= 0.75\n",
    "GREEN_IMAGE[2, :, :] *= 0.25\n",
    "GREEN_IMAGE = nn.Parameter(GREEN_IMAGE, requires_grad=True)\n",
    "imshow(BLUE_IMAGE.detach())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa79d1e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "imshow(GREEN_IMAGE.detach())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e633d600",
   "metadata": {},
   "outputs": [],
   "source": [
    "#           X                            Y\n",
    "train_dataset = [\n",
    "    # Good things, go forward to little green frogs and little blue birds\n",
    "    [['little', 'green', 'frog'],                  FORWARD],\n",
    "    [['little', GREEN_IMAGE, 'frog'],              FORWARD],\n",
    "\n",
    "    # Bad things, go backward to big green hulks and big blue oceans\n",
    "    [['big', 'blue', 'ocean', BLUE_IMAGE],       BACKWARD],\n",
    "    [['blue', 'big', BLUE_IMAGE, 'ocean'],       BACKWARD],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c58fcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "atari_decoding = {\n",
    "    0: 'FORWARD',\n",
    "    1: 'BACKWARD',\n",
    "    2: 'LEFT',\n",
    "    3: 'RIGHT',\n",
    "}\n",
    "\n",
    "for key in embedding_layer:\n",
    "    if not isinstance(embedding_layer[key], nn.Parameter):\n",
    "        embedding_layer[key] = nn.Parameter(embedding_layer[key], requires_grad=True)\n",
    "\n",
    "def embed(sample):\n",
    "    embeddings = None\n",
    "    for token in sample:\n",
    "        if isinstance(token, str):\n",
    "            if embeddings is None:\n",
    "                embeddings = embedding_layer[token].unsqueeze(0)\n",
    "            else:\n",
    "                embeddings = torch.cat([embeddings, embedding_layer[token].unsqueeze(0)])\n",
    "        elif isinstance(token, torch.Tensor):\n",
    "            if embeddings is None:\n",
    "                if token.shape == (3, 224, 224):\n",
    "                    embeddings = example_image_embedding(token).unsqueeze(0)\n",
    "                else:\n",
    "                    embeddings = atari_embedding(token).unsqueeze(0)\n",
    "            else:\n",
    "                if token.shape == (3, 224, 224):\n",
    "                    embeddings = torch.cat([embeddings, example_image_embedding(token).unsqueeze(0)])\n",
    "                else:\n",
    "                    embeddings = torch.cat([embeddings, atari_embedding(token).unsqueeze(0)])\n",
    "        else: \n",
    "            raise ValueError(f\"Unknown token type: {type(token)}\")\n",
    "    return embeddings\n",
    "\n",
    "def decode(token):\n",
    "    if token >= len(vocabulary):\n",
    "        return atari_decoding[(token - len(embedding_layer)).item()]\n",
    "    else:\n",
    "        return vocabulary[token]\n",
    "\n",
    "def encode(target):\n",
    "    if isinstance(target, str):\n",
    "        return tensor(vocab_index_of(target))\n",
    "    else:\n",
    "        return target + len(embedding_layer)\n",
    "\n",
    "def combine_all_embeddings(embeddings):\n",
    "    return torch.mean(embeddings, dim=0)\n",
    "\n",
    "def vocab_index_of(token):\n",
    "    return next((i for i, t in enumerate(embedding_layer) if t == token), -1)\n",
    "\n",
    "def loss_fn(prediction, target):\n",
    "    if isinstance(target, str):\n",
    "        target = tensor(vocab_index_of(target))\n",
    "    else:\n",
    "        target = target + len(embedding_layer)\n",
    "    return nn.functional.cross_entropy(prediction, target)\n",
    "\n",
    "num_out = (\n",
    "    len(embedding_layer)   # number of words in our vocabulary\n",
    "    + len(atari_embedding.weight) # number of button presses\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "955f4e70",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(3, 2),\n",
    "    nn.Sigmoid(),\n",
    "    nn.Linear(2, num_out),\n",
    ")\n",
    "optimizer = torch.optim.AdamW([\n",
    "    {'params': model.parameters()},\n",
    "    {'params': embedding_layer.values()},\n",
    "    {'params': [BLUE_IMAGE, GREEN_IMAGE]},\n",
    "    {'params': atari_embedding.parameters()}\n",
    "    ], lr=0.001)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c88c300c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample, target = ['little', GREEN_IMAGE, 'frog'], FORWARD  # Expect \"little\" and GREEN IMAGE to be close to something like \"frog\" or \"green\" or FORWARD. (Little green frogs are good.)\n",
    "\n",
    "# Forward to make a prediction.\n",
    "embeddings = combine_all_embeddings(embed(sample))\n",
    "\n",
    "predictions = model(embeddings)\n",
    "predicted_token = decode(predictions.argmax())\n",
    "\n",
    "print(f'Expected: {decode(target + len(vocabulary))}\\nGot: {predicted_token}\\nWith confidence {predictions.max():.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6d3d7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample, target = ['little', GREEN_IMAGE, 'frog'], FORWARD\n",
    "embeddings = combine_all_embeddings(embed(sample))\n",
    "prediction = model(embeddings)\n",
    "loss = nn.functional.cross_entropy(prediction, encode(target))\n",
    "loss.backward()\n",
    "optimizer.step()\n",
    "optimizer.zero_grad()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6888f584",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample, target = ['little', GREEN_IMAGE, 'frog'], FORWARD\n",
    "embeddings = combine_all_embeddings(embed(sample))\n",
    "predictions = model(embeddings)\n",
    "predicted_token = decode(predictions.argmax())\n",
    "\n",
    "print(f'Expected: {target}\\nGot: {predicted_token}\\nWith confidence {predictions.max():.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "180df9db",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_layer['blue'], model(embeddings), list(model.parameters())[0], model(embeddings).max(), model(embeddings).argmax()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b37335f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample, target = train_dataset[0]\n",
    "embeddings = combine_all_embeddings(embed(sample))\n",
    "prediction = model(embeddings)\n",
    "optimizer.zero_grad()\n",
    "loss = nn.functional.cross_entropy(prediction, encode(target))\n",
    "loss.backward()\n",
    "optimizer.step()\n",
    "target, tensor(vocab_index_of(target)), vocabulary[4], prediction, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4191a998",
   "metadata": {},
   "outputs": [],
   "source": [
    "target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d36ba46",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_layer['blue'], model(embeddings), list(model.parameters())[0], model(embeddings).max(), model(embeddings).argmax()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9337b3e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math; -math.log(1/12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ae2964a",
   "metadata": {},
   "outputs": [],
   "source": [
    "atari_embedding(FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27eee37c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    example_embedding_1('big'),\n",
    "    example_embedding_1('blue'),\n",
    "    example_embedding_1('ocean'),\n",
    "    example_image_embedding(BLUE_IMAGE),\n",
    "    atari_embedding(BACKWARD),\n",
    "    example_embedding_1('little'),\n",
    "    example_embedding_1('green'),\n",
    "    example_embedding_1('frog'),\n",
    "    example_image_embedding(GREEN_IMAGE),\n",
    "    atari_embedding(FORWARD),\n",
    "), colors=[\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f45c268c",
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa6b522",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8151ea3",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for _ in tqdm(range(1000)):\n",
    "    for sample, target in train_dataset:\n",
    "        embeddings = combine_all_embeddings(embed(sample))\n",
    "        probabilities = model(embeddings)\n",
    "        loss = loss_fn(probabilities, target)\n",
    "        losses.append(loss.item())\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45c8d3a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    example_embedding_1('big'),\n",
    "    example_embedding_1('blue'),\n",
    "    example_embedding_1('ocean'),\n",
    "    example_image_embedding(BLUE_IMAGE),\n",
    "    atari_embedding(BACKWARD),\n",
    "    example_embedding_1('little'),\n",
    "    example_embedding_1('green'),\n",
    "    example_embedding_1('frog'),\n",
    "    example_image_embedding(GREEN_IMAGE),\n",
    "    atari_embedding(FORWARD),\n",
    "), colors=[\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "    'green',\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c48931e",
   "metadata": {},
   "outputs": [],
   "source": [
    "atari_embedding(BACKWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43182b4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot()\n",
    "ax.plot(losses)\n",
    "# Display the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19b3577d",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample, target in train_dataset:\n",
    "    embeddings = embed(sample)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1822b3b-1719-4704-a931-1b1614ec5c9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiktoken_tokenizer.encode('cloudy sky')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50ca2539-5abd-4d12-8682-831bd9c97ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'{\"token\":>10}: {\"text\":>10}')\n",
    "\n",
    "for token in tiktoken_tokenizer.encode(\"cloudy sky\"):\n",
    "    decoded = tiktoken_tokenizer.decode([token])\n",
    "    print(f'{token:>10}: {decoded:>10}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a7edd1a-aa97-4dd7-92ec-e47fd02a0791",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'{\"token\":>10}: {\"text\":>10}')\n",
    "\n",
    "for token in tiktoken_tokenizer.encode(\"Hello, abcDEF world!\"):\n",
    "    decoded = tiktoken_tokenizer.decode([token])\n",
    "    print(f'{token:>10}: {decoded:>10}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55491ef6-1f75-43d4-8c7a-8692999fa4bd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7eb58a-56fd-49d6-a88b-e724e92cca94",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = spaces_tokenizer('a lake')\n",
    "embeddings = [\n",
    "    example_embedding_2(token) for token in tokens\n",
    "]\n",
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b82e5be5-005e-41d5-b883-542aa1c34d2d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "84ec9970-f97f-4780-bc9d-5ccadca49f3a",
   "metadata": {},
   "source": [
    "## Tokenizing images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b084cb60-6eb7-4d55-aa70-4d3bf0958450",
   "metadata": {},
   "source": [
    "#### Visual question answering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f4212f9-7170-480f-8ea5-f02096690275",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_index = 10\n",
    "sample = vqa_dataset['train'][sample_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb823d8-0f3f-4f23-ad1a-d386111e8893",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3173ed29-a85f-4e8d-a296-f8fc01a0fe69",
   "metadata": {},
   "source": [
    "We want to encode this sample as:\n",
    "\n",
    "```\n",
    "[question...,  image...,  answer...   ]\n",
    "     ^           ^           ^\n",
    "     |           |           |\n",
    "[  given this prefix   ,  predict this]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04c40b82-ab4f-470d-b4a7-44aaa9f8a99e",
   "metadata": {},
   "source": [
    "#### Robotic movement and control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e4fe4d2-feee-4d8a-a8be-7fce12b97ce3",
   "metadata": {},
   "outputs": [],
   "source": [
    "robotics_sample = four_rooms_dataset['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58e1e1e6-3367-4f8f-9ada-dda2efb4988e",
   "metadata": {},
   "outputs": [],
   "source": [
    "robotics_sample.observations.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33bddc46-5eba-466e-a342-8db5e3fc7cd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "robotics_sample.actions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e0c00e4-0060-4859-b29d-eaceae9681e1",
   "metadata": {},
   "source": [
    "We want to encode this sample as:\n",
    "\n",
    "```\n",
    "[mission..., image..., direction...,   action...  ]\n",
    "   ^           ^            ^             ^\n",
    "   |           |            |             |\n",
    "[      given this prefix...        , predict this ]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b08d97a3-a4f0-4e59-aed4-323c040900d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample['question']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "065aacc1-4ba9-427c-b692-b1179a574df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "spaces_tokenizer(sample['question'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4daba7a-4882-4af1-b923-5b07abf9199a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.encode_text(sample['question'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9efeecda-1638-43c6-bb7b-8e5986ab3b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample['answers'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ad6fa56-71d1-4bb0-88b2-c8b7044bc3ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample['image']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fc1fc0e-70d0-41dc-b2b8-b86ab5c0e69f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e6a9bdc-078a-4b7d-967f-1fb4f1d77522",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b5c914-ab6c-4e9f-ad1d-a7b636dd19fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs, labels, attention_mask = batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec19000-ccdc-4c3c-90c3-cb9493a975c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs['question'][2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7093a9a4-bb47-4519-ba7e-7e8b7f445c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.decode_text(inputs['question'][2].squeeze())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c63c8bb6-4570-4ac5-a0b9-d9d2ddc28f93",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65268629-1abe-4c92-9fb8-8938395f3013",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e477860c-08b4-4385-99c0-26ea851d29bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "412865a2-170a-437c-b3b7-7f491eb38f7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_to_token = {\n",
    "    'pond': 0,\n",
    "    'ocean': 1,\n",
    "    'frog': 2,\n",
    "    'hulk': 3,\n",
    "    'blue': 4,\n",
    "    'green': 5,\n",
    "    'big': 6,\n",
    "    'little': 7,\n",
    "}\n",
    "vocab_size = len(vocab_to_token)\n",
    "\n",
    "token_to_vocab = {\n",
    "    0: 'pond',\n",
    "    1: 'ocean',\n",
    "    2: 'frog',\n",
    "    3: 'hulk',\n",
    "    4: 'blue',\n",
    "    5: 'green',\n",
    "    6: 'big',\n",
    "    7: 'little',\n",
    "}\n",
    "\n",
    "embedding_layer = nn.Embedding(8, 3)\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(3, 16),\n",
    "    nn.Linear(16, vocab_size),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24856cc7-0361-4493-8376-6d39d492d2e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding = embedding_layer(tensor([vocab_to_token[token] for token in spaces_tokenizer('little blue pond')]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b80763cb-8ce0-4823-ba1b-fc294ab11f94",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04ce56f5-326c-4404-a37b-8bc73d503a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "model(embedding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75189851-f797-49ed-9f87-7e7f1d9864f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "blue_image = torch.ones(3, 224, 224)\n",
    "blue_image[0, :, :] *= 0.25\n",
    "blue_image[1, :, :] *= 0.25\n",
    "blue_image[2, :, :] *= 0.75\n",
    "green_image = torch.ones(3, 224, 224)\n",
    "green_image[0, :, :] *= 0.25\n",
    "green_image[1, :, :] *= 0.75\n",
    "green_image[2, :, :] *= 0.25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5343f3d-04c8-49e9-97b1-cba2701252fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_embedding(*zip(\n",
    "    embedding(tensor(vocab_to_token['pond'])),\n",
    "    embedding(tensor(vocab_to_token['ocean'])),\n",
    "    embedding(tensor(vocab_to_token['frog'])),\n",
    "    embedding(tensor(vocab_to_token['hulk'])),\n",
    "    embedding(tensor(vocab_to_token['blue'])),\n",
    "    embedding(tensor(vocab_to_token['green'])),\n",
    "    embedding(tensor(vocab_to_token['big'])),\n",
    "    embedding(tensor(vocab_to_token['little'])),\n",
    "    example_image_embedding(blue_image),\n",
    "    example_image_embedding(green_image),\n",
    "), colors=[\n",
    "    'blue',\n",
    "    'blue',\n",
    "    'green',\n",
    "    'green',\n",
    "    'blue',\n",
    "    'green',\n",
    "    'red',\n",
    "    'red',\n",
    "    'lightblue',\n",
    "    'lightgreen',\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd81623-b042-48eb-86a9-f7b1d4c8b359",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = [\n",
    "    'big',\n",
    "    'blue',\n",
    "    'ocean',\n",
    "    'big blue',\n",
    "    'big ocean',\n",
    "    'blue ocean',\n",
    "    'big blue ocean',\n",
    "\n",
    "    'little',\n",
    "    'green',\n",
    "    'little green',\n",
    "    'little frog',\n",
    "    'green frog',\n",
    "    'little green frog',\n",
    "\n",
    "    'hulk',\n",
    "    'big hulk',\n",
    "    'green hulk',\n",
    "    # 'big green hulk',\n",
    "\n",
    "    'pond',\n",
    "    'little pond',\n",
    "    'blue pond',\n",
    "    # 'little blue pond',\n",
    "]\n",
    "\n",
    "val_dataset = [\n",
    "    'little blue pond',\n",
    "    'big green hulk',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f69df6cd-8bb3-46ab-b312-8c574d850df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample in train_dataset:\n",
    "    tokens = spaces_tokenizer(sample)\n",
    "    print(tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da5b347e-7c96-45da-b9cb-7f0b9fc7ebee",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample in train_dataset:\n",
    "    tokens = [vocab_to_token[x] for x in spaces_tokenizer(sample)]\n",
    "    print(tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b4d055-9796-407d-a00d-ec8c7d762ca6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0c53f07-89c1-4f32-acae-9061ec094925",
   "metadata": {},
   "outputs": [],
   "source": [
    "for sample in train_dataset:\n",
    "    tokens = [vocab_to_token[x] for x in spaces_tokenizer(sample)]\n",
    "    embeddings = embedding(tensor(tokens))\n",
    "    print(embeddings)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}