path: root/transformer_shortest_paths.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LPphBnKR-aWF"
      },
      "source": [
        "# Step 0: Imports"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "I need to do something simpler first.\n",
        "\n",
        "\n",
        "1. Train a transformer to output 1 if token x is in the input, and 0 else.\n",
        "\n",
        "2. Train a transformer to output 1 if token x and token y are both in the input and 0 else.\n",
        "\n",
        "3. Train a transformer to output 1 if token x and token y are adjacent in the input.\n",
        "\n",
        "4. Train a transformer to output 1 if token x and token y are adjacent in the input AND they're 2k, 2k+1\n",
        "\n"
      ],
      "metadata": {
        "id": "HscaSHV43vU0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ge5QvElvhCOw",
        "outputId": "c7cdaefa-d6dc-44ad-c258-e4fb2aca97a5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "imports complete\n"
          ]
        }
      ],
      "source": [
        "# imports\n",
        "import numpy as np\n",
        "from collections import deque\n",
        "import pickle\n",
        "from tqdm import tqdm\n",
        "np.random.seed(42)\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import pickle\n",
        "from math import sqrt\n",
        "from torch.utils.data import DataLoader, TensorDataset\n",
        "import matplotlib.pyplot as plt\n",
        "torch.manual_seed(42)\n",
        "\n",
        "import os\n",
        "\n",
        "print(\"imports complete\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lylOX2POPwFL"
      },
      "outputs": [],
      "source": [
        "SEQ_LEN = 32\n",
        "\n",
        "PAD_TOKEN = 0\n",
        "AVG_DEG = 2\n",
        "MAX_VTXS = SEQ_LEN//AVG_DEG - 1\n",
        "# vertices are labelled 1,2,...,63\n",
        "# we also have a padding token which is 0.\n",
        "\n",
        "INF = MAX_VTXS # represents unreachability"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKt-yIpDebF1"
      },
      "source": [
        "# Step 1: Generate synthetic data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1IbzGIWseK3E",
        "outputId": "a3cbc233-358c-4e17-ea6e-f4e9349d886b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 1/1 [00:14<00:00, 14.42s/it]\n"
          ]
        }
      ],
      "source": [
        "# original task data\n",
        "NTRAIN1 = 100_000\n",
        "# the data will be edge lists\n",
        "# like this: [1 3 1 5 2 4 0 0 0 0]\n",
        "# this represents edges (1,3), (1,5) (2,4)\n",
        "# (the zeros are just padding tokens)\n",
        "\n",
        "# the label is the shortest distance from vtx 1 to vtx 2\n",
        "# or \"INF\" if no path exists\n",
        "\n",
        "# fine tuning data\n",
        "NTRAIN2 = 2000\n",
        "# I haven't totally figured out how to do the fine tuning yet.\n",
        "# So don't worry about this yet.\n",
        "\n",
        "def random_graph(n):\n",
        "    edge_list = []\n",
        "    adjacencies = [set() for _ in range(n+1)]\n",
        "    indices = np.random.randint(n, size=(AVG_DEG*(n-1)))+1\n",
        "    for i in range(0, len(indices), 2):\n",
        "        u = indices[i]\n",
        "        v = indices[i + 1]\n",
        "        if u != v:\n",
        "            edge_list += [u,v]\n",
        "            adjacencies[u].add(v)\n",
        "            adjacencies[v].add(u)\n",
        "\n",
        "    if np.random.random() < 0.25:\n",
        "      edge_list += [1,2]\n",
        "      adjacencies[1].add(2)\n",
        "      adjacencies[2].add(1)\n",
        "\n",
        "    edge_list += [PAD_TOKEN]*(SEQ_LEN-len(edge_list))\n",
        "    return edge_list, adjacencies\n",
        "\n",
        "\"\"\"\n",
        "input: G, represented as an adjacency list\n",
        "output: [INF]+[d(1,i) for i in range(n)] if target=None\n",
        "if target is set to some value, then we instead just output that specific distance\n",
        "\"\"\"\n",
        "def SSSP(G, target=None):\n",
        "    dist = [INF for _ in G]\n",
        "    dist[1] = 0\n",
        "    frontier = deque()\n",
        "    frontier.append(1)\n",
        "    while len(frontier) > 0:\n",
        "        vtx = frontier.popleft()\n",
        "        for x in G[vtx]:\n",
        "            if dist[x] == INF:\n",
        "                dist[x] = 1 + dist[vtx]\n",
        "                frontier.append(x)\n",
        "                if x == target:\n",
        "                    return dist[target]\n",
        "    if target is not None:\n",
        "        return dist[target]\n",
        "    else:\n",
        "        return dist\n",
        "\n",
        "def fake_SSSP(G, target=None):\n",
        "    return 2 in G[1]\n",
        "\n",
        "graphs1 = []\n",
        "distance1 = []\n",
        "\n",
        "graphs2 = []\n",
        "distances2 = []\n",
        "\n",
        "for n in tqdm(range(MAX_VTXS-1, MAX_VTXS)):\n",
        "    # for _ in range(NTRAIN1//MAX_VTXS):\n",
        "    for _ in range(NTRAIN1):\n",
        "        edge_list, adj_list = random_graph(n)\n",
        "        dist = SSSP(adj_list, target=2)\n",
        "\n",
        "        graphs1.append(edge_list)\n",
        "        distance1.append(dist)\n",
        "\n",
        "# for n in range(8, MAX_VTXS//4):\n",
        "#     for _ in range(NTRAIN2//MAX_VTXS):\n",
        "#         edge_list, adj_list = random_graph(n)\n",
        "#         distances = SSSP(adj_list)\n",
        "#         graphs2.append(edge_list)\n",
        "#         distances2.append(distances)\n",
        "\n",
        "split1 = int(len(graphs1)*3/4)\n",
        "split2 = int(len(graphs2)*3/4)\n",
        "\n",
        "all1 = list(zip(graphs1, distance1))\n",
        "np.random.shuffle(all1)\n",
        "graphs1, distance1 = zip(*all1)\n",
        "\n",
        "data = {\n",
        "    \"train1-data\": graphs1[:split1],\n",
        "    \"train1-labels\": distance1[:split1],\n",
        "    \"test1-data\": graphs1[split1:],\n",
        "    \"test1-labels\": distance1[split1:]\n",
        "    # \"train2-data\": graphs2[:split2],\n",
        "    # \"train2-labels\": distances2[:split2],\n",
        "    # \"test2-data\": graphs2[split2:],\n",
        "    # \"test2-labels\": distances2[split2:]\n",
        "}\n",
        "\n",
        "with open('data.pkl', 'wb') as file:\n",
        "    pickle.dump(data, file)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EpDBxcgaIPpJ",
        "outputId": "37cf9577-8cd8-444c-ec1a-c6f4b6061b7f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dataset size = 45MB\n"
          ]
        }
      ],
      "source": [
        "print(f\"dataset size = {os.path.getsize('data.pkl')//(1024*1024)}MB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q3Cg_8UQep8g"
      },
      "source": [
        "# Step 2: Define Transformer Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tLOWhg_CeWzH"
      },
      "outputs": [],
      "source": [
        "class TransformerModel(nn.Module):\n",
        "    def __init__(self, input_dim, model_dim, output_dim, num_heads, num_layers, seq_len, device, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.embedding = nn.Embedding(input_dim, model_dim)\n",
        "        self.model_dim = model_dim\n",
        "        self.seq_len = seq_len\n",
        "        self.device = device\n",
        "\n",
        "        # weight sharing\n",
        "        encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=num_heads,\n",
        "                                                   dim_feedforward=model_dim*4,\n",
        "                                                   dropout=dropout, batch_first=True)\n",
        "        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)\n",
        "\n",
        "        self.fc_out = nn.Linear(model_dim*seq_len, output_dim)\n",
        "\n",
        "    # def positional_encoding(self, batch_size):\n",
        "    #     pos_encoding = torch.arange(self.seq_len, device=self.device).unsqueeze(1)\n",
        "    #     pos_encoding = pos_encoding.float().unsqueeze(0).repeat(batch_size, 1, 1)\n",
        "    #     return pos_encoding\n",
        "\n",
        "    def positional_encoding(self, batch_size):\n",
        "        position = torch.arange(self.seq_len, dtype=torch.float, device=self.device).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, self.model_dim, 2, dtype=torch.float, device=self.device) *\n",
        "                            -(torch.log(torch.tensor(500.0)) / self.model_dim))\n",
        "\n",
        "        pos_encoding = torch.zeros(self.seq_len, self.model_dim, device=self.device)\n",
        "        pos_encoding[:, 0::2] = torch.sin(position * div_term)\n",
        "        pos_encoding[:, 1::2] = torch.cos(position * div_term)\n",
        "        pos_encoding = pos_encoding.unsqueeze(0).repeat(batch_size, 1, 1)\n",
        "        return pos_encoding\n",
        "\n",
        "    def forward(self, src, key_padding_mask):\n",
        "        batch_size, src_len = src.size(0), src.size(1)\n",
        "        src_pos = self.positional_encoding(batch_size)\n",
        "        embed = self.embedding(src)\n",
        "        src = embed * sqrt(self.model_dim) + src_pos\n",
        "\n",
        "        output = self.transformer_encoder(src, None, src_key_padding_mask=key_padding_mask)\n",
        "        flat_output = torch.flatten(output, start_dim=1, end_dim=2)\n",
        "        output = self.fc_out(flat_output)\n",
        "        return output\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bpIeg86S-hBb"
      },
      "source": [
        "# Step 3: Load Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kWXvJRDYgFVP",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c13adb9d-6565-43b5-8437-20cef3dc0d16"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Trainable parameters in the model: 26K\n",
            "train BASELINEs: 39.4069\n"
          ]
        }
      ],
      "source": [
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "assert device.type == 'cuda', \"CUDA is not available. Please check your GPU setup.\"\n",
        "\n",
        "# PARAMS\n",
        "VOCAB_SIZE = 1+MAX_VTXS # one more than the max number of vertices\n",
        "MODEL_DIM = 32 # Dimension of model (embedding and transformer)\n",
        "NEPOCHS = 10\n",
        "BSZ = 32\n",
        "LR = 0.01\n",
        "NHEADS = 4\n",
        "NLAYERS = 2\n",
        "PAD_TOKEN = 0\n",
        "model = TransformerModel(input_dim=VOCAB_SIZE, model_dim=MODEL_DIM,\n",
        "                         output_dim=1, num_heads=NHEADS,\n",
        "                         num_layers=NLAYERS, seq_len=SEQ_LEN,\n",
        "                         device=device).to(device)\n",
        "\n",
        "with open(\"data.pkl\", \"rb\") as f:\n",
        "    data = pickle.load(f)\n",
        "\n",
        "train_data1 = data[\"train1-data\"]\n",
        "train_label1 = data[\"train1-labels\"]\n",
        "train_data_tensor = torch.tensor(train_data1, dtype=torch.long, device=device)\n",
        "train_label_tensor = torch.tensor(train_label1, dtype=torch.float, device=device)\n",
        "train_padding_mask = (train_data_tensor != PAD_TOKEN).bool().to(device)\n",
        "train_dataset = TensorDataset(train_data_tensor, train_label_tensor, train_padding_mask)\n",
        "train_loader = DataLoader(train_dataset, batch_size=BSZ, shuffle=True)\n",
        "\n",
        "test_data1 = data[\"test1-data\"]\n",
        "test_label1 = data[\"test1-labels\"]\n",
        "test_data_tensor = torch.tensor(test_data1, dtype=torch.long, device=device)\n",
        "test_label_tensor = torch.tensor(test_label1, dtype=torch.float, device=device)\n",
        "test_padding_mask = (test_data_tensor != PAD_TOKEN).bool().to(device)\n",
        "test_dataset = TensorDataset(test_data_tensor, test_label_tensor, test_padding_mask)\n",
        "test_loader = DataLoader(test_dataset, batch_size=BSZ, shuffle=True)\n",
        "\n",
        "criterion = nn.MSELoss()\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=LR)\n",
        "\n",
        "train_err = []\n",
        "test_err = []\n",
        "\n",
        "trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "print(f\"Trainable parameters in the model: {trainable_params//1000}K\")\n",
        "\n",
        "train_baseline = ((train_label_tensor - train_label_tensor.mean())**2).mean().item()\n",
        "print(f\"train BASELINEs: {train_baseline:.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "205MvfJQQYya"
      },
      "source": [
        "# Dad reccomended having more \"partial progress measures\" / having the model \"show it's work\".\n",
        "# or creating a different easier training regiment to start."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f8Zn33m7CxL5"
      },
      "source": [
        "# Step 4: Train the Model for the first task"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 486
        },
        "id": "pvTfzGmCeXU4",
        "outputId": "0d3a20f3-23be-4c19-9eb6-46bfe11a48b1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/10 \t Train Err: 28.8616 \t Test Err: 39.4354 \t baseline err: 39.4069\n",
            "Epoch 2/10 \t Train Err: 39.8088 \t Test Err: 39.4255 \t baseline err: 39.4069\n",
            "Epoch 3/10 \t Train Err: 39.7257 \t Test Err: 39.9765 \t baseline err: 39.4069\n",
            "Epoch 4/10 \t Train Err: 39.4951 \t Test Err: 40.0988 \t baseline err: 39.4069\n",
            "Epoch 5/10 \t Train Err: 39.4205 \t Test Err: 39.5148 \t baseline err: 39.4069\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "KeyboardInterrupt",
          "evalue": "",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-17-3dc1bf4cf066>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      8\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcriterion\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_labels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m         \u001b[0mtrain_loss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_loader\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 10\u001b[0;31m         \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    519\u001b[0m                 \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    520\u001b[0m             )\n\u001b[0;32m--> 521\u001b[0;31m         torch.autograd.backward(\n\u001b[0m\u001b[1;32m    522\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    523\u001b[0m         )\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    287\u001b[0m     \u001b[0;31m# some Python versions print out the first line of a multi-line function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    288\u001b[0m     \u001b[0;31m# calls in the traceback and some print out the last line\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 289\u001b[0;31m     _engine_run_backward(\n\u001b[0m\u001b[1;32m    290\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    291\u001b[0m         \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/autograd/graph.py\u001b[0m in \u001b[0;36m_engine_run_backward\u001b[0;34m(t_outputs, *args, **kwargs)\u001b[0m\n\u001b[1;32m    767\u001b[0m         \u001b[0munregister_hooks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_register_logging_hooks_on_whole_graph\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt_outputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    768\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 769\u001b[0;31m         return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[1;32m    770\u001b[0m             \u001b[0mt_outputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    771\u001b[0m         )  # Calls into the C++ engine to run the backward pass\n",
            "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
          ]
        }
      ],
      "source": [
        "for epoch in range(NEPOCHS):\n",
        "    model.train() # set to training mode\n",
        "    train_loss = 0\n",
        "\n",
        "    for batch_src, batch_labels, batch_padding_mask in train_loader:\n",
        "        optimizer.zero_grad()\n",
        "        output = model(batch_src, batch_padding_mask)\n",
        "        loss = criterion(output.squeeze(1), batch_labels)\n",
        "        train_loss += loss.item()/len(train_loader)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    # Evaluate performance\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch_src, batch_labels, batch_padding_mask in test_loader:\n",
        "            output = model(batch_src, batch_padding_mask)\n",
        "            loss = criterion(output.squeeze(1), batch_labels)\n",
        "            test_loss += loss.item()/len(test_loader)\n",
        "\n",
        "    test_err.append(test_loss)\n",
        "    train_err.append(train_loss)\n",
        "    print(f\"Epoch {epoch + 1}/{NEPOCHS} \\t Train Err: {train_loss:.4f} \\t Test Err: {test_loss:.4f} \\t baseline err: {train_baseline:.4f}\")\n",
        "\n",
        "plt.figure(figsize=(10, 5))\n",
        "plt.plot(test_err, label='Test', color='red')\n",
        "plt.plot(train_err, label='Train', color='blue')\n",
        "plt.title('Accuracy vs Epochs')\n",
        "plt.xlabel('Epochs'); plt.ylabel('Accuracy')\n",
        "plt.legend(); plt.grid()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v1hCiItHDWxJ"
      },
      "outputs": [],
      "source": [
        "## Q: why is this not working so well?\n",
        "\n",
        "## maybe first try a simpler problem: just give it points for distinguishing between distance 1 or not"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LoGEmM5lH7_A"
      },
      "outputs": [],
      "source": [
        "batch_src, batch_labels, batch_padding_mask = next(iter(train_loader))\n",
        "output = model(batch_src, batch_padding_mask)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hO8AhX3G7vF8",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8f4a3ca6-db47-434d-95a4-4631bc73de62"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1 \t 5.7\n",
            "3 \t 5.7\n",
            "15 \t 7.1\n",
            "1 \t 5.7\n",
            "2 \t 7.8\n",
            "15 \t 7.1\n",
            "1 \t 0.7\n",
            "15 \t 5.7\n",
            "5 \t 5.7\n",
            "1 \t 5.7\n",
            "1 \t 0.7\n",
            "4 \t 5.7\n",
            "2 \t 7.8\n",
            "3 \t 5.7\n",
            "3 \t 5.7\n",
            "15 \t 7.8\n",
            "15 \t 7.8\n",
            "1 \t 5.7\n",
            "3 \t 7.1\n",
            "1 \t 5.7\n",
            "3 \t 5.7\n",
            "1 \t 7.1\n",
            "1 \t 7.8\n",
            "2 \t 5.7\n",
            "1 \t 5.7\n",
            "15 \t 7.1\n",
            "6 \t 7.1\n",
            "1 \t 5.7\n",
            "1 \t 5.7\n",
            "1 \t 5.7\n",
            "15 \t 7.1\n",
            "1 \t 7.1\n"
          ]
        }
      ],
      "source": [
        "for x,y in zip(batch_labels.tolist(),  output.squeeze(1).tolist()):\n",
        "  print(f\"{int(x)} \\t {y:.1f}\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "batch_src[2]"
      ],
      "metadata": {
        "id": "dRdUGbFmkPtK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LC6Xv3YfC0Rm"
      },
      "source": [
        "# Step 5: Fine Tune"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JtTLXn4zC1z_"
      },
      "source": [
        "# Step 6: Test generalization"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}