{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "34a31b9f-e4e1-47cd-a214-c2b4a3a12050",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16640b2a-025f-441c-9ae2-f38222d15005",
   "metadata": {},
   "source": [
    "### Tutorial A2: Predictions\n",
    "\n",
    "When one has a machine learning model, a natural next step is to make predictions with that model. Accordingly, `tangermeme` implements general-purpose functions for making predictions from PyTorch models in a memory-efficient manner, regardless of the number of inputs or outputs from the model. These functions can be used by themselves but their primary purpose is as building blocks for more complicated analysis functions. \n",
    "\n",
    "Although making predictions using a model is conceptually simple to understand, there are several technical issues with doing efficiently so in practice. First, when your data is too big to fit in GPU memory you cannot simply move it all over at once and so you must make predictions in a batched manner. Here, a small number of examples are moved from the CPU to the GPU, predictions are made, and the results are moved back to the CPU. The batch size can be tuned to the largest number of examples that fit in GPU memory. Second, when a model has multiple outputs, making predictions in a batched manner yields a set of tensors for each batch. These tensors must be correctly concatenated across batches to make sure that the final output from the function matches the shape of the data as if all examples were run through the model at the same time. Third, some models have multiple inputs and so this function must be able to handle an optional set of additional arguments."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f99b808d-7c73-48a3-8b64-dbc548bb75e3",
   "metadata": {},
   "source": [
    "#### Predict\n",
    "\n",
    "The simplest function that implements these ideas is `predict`, which takes in a model, data, and optional additional arguments, and makes batched predictions on the data given the model. The function can be run on GPUs, CPUs, or any other devices that work with PyTorch.\n",
    "\n",
    "To demonstate it, let's use a model that takes its inputs and flattens them before feeding them into a dense layer to make three predictions per example. The forward function takes in two optional arguments: `alpha`, which gets added to the predictions, and `beta`, which multiplies the predictions (but not `alpha`). By default, these are set such that the predictions are returned without modification."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ac89068e-1aaf-4d8b-a4ef-0b21690ed005",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "class FlattenDense(torch.nn.Module):\n",
    "\tdef __init__(self, length=10):\n",
    "\t\tsuper(FlattenDense, self).__init__()\n",
    "\t\tself.dense = torch.nn.Linear(length*4, 3)\n",
    "\n",
    "\tdef forward(self, X, alpha=0, beta=1):\n",
    "\t\tX = X.reshape(X.shape[0], -1)\n",
    "\t\treturn self.dense(X) * beta + alpha"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af9607f4-5fbf-45fe-8147-9aedf5e20428",
   "metadata": {},
   "source": [
    "Now let's generate some random sequence and see what the model output looks like for it. We need to use `torch.manual_seed` because of the random initializations used in the `torch.nn.Linear` layer. As a side note, even though we are not training the model here, the usage doesn't change based on whether the model is randomly initialized or trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "564f94e0-84e0-4dcb-9de7-d32bb1d3bf77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3154, -0.1625, -0.3183],\n",
       "        [-0.0866,  0.5461, -0.0244],\n",
       "        [ 0.3089, -0.2828, -0.1485],\n",
       "        [ 0.1671, -0.1341, -0.3094],\n",
       "        [-0.0627,  0.0088,  0.3471]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tangermeme.utils import random_one_hot\n",
    "torch.manual_seed(0)\n",
    "\n",
    "X = random_one_hot((5, 4, 10), random_state=0).float()\n",
    "model = FlattenDense()\n",
    "\n",
    "y = model(X)\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12667e15-6bbc-499b-be79-21441cbdd112",
   "metadata": {},
   "source": [
    "This is simple enough to do for a simple model on a small amount of data. Let's try using the built-in predict function with different batch sizes, to demonstrate how one would do batched predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1870e566-c1f5-4bd5-b500-587e3867ce6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3154, -0.1625, -0.3183],\n",
       "        [-0.0866,  0.5461, -0.0244],\n",
       "        [ 0.3089, -0.2828, -0.1485],\n",
       "        [ 0.1671, -0.1341, -0.3094],\n",
       "        [-0.0627,  0.0088,  0.3471]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tangermeme.predict import predict\n",
    "\n",
    "y0 = predict(model, X, batch_size=2)\n",
    "y0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cf127ba7-91fe-428d-aab9-23bbc4d1806e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3154, -0.1625, -0.3183],\n",
       "        [-0.0866,  0.5461, -0.0244],\n",
       "        [ 0.3089, -0.2828, -0.1485],\n",
       "        [ 0.1671, -0.1341, -0.3094],\n",
       "        [-0.0627,  0.0088,  0.3471]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y0 = predict(model, X, batch_size=100)\n",
    "y0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62557ffe-e5b6-4d29-ba05-d5e06e002433",
   "metadata": {},
   "source": [
    "Note that the tensor no longer has `grad_fn=<AddBackward0>` meaning that gradients were not calculated or stored. Specifically, the prediction loop is wrapped in `torch.no_grad`. By default, this function will move each batch to the GPU. However, it doesn't have to. You can pass `device='cpu'` to have the predictions be made on the CPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "746c4787-99f7-422f-aa9d-35c5e3ca6e8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3154, -0.1625, -0.3183],\n",
       "        [-0.0866,  0.5461, -0.0244],\n",
       "        [ 0.3089, -0.2828, -0.1485],\n",
       "        [ 0.1671, -0.1341, -0.3094],\n",
       "        [-0.0627,  0.0088,  0.3471]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y0 = predict(model, X, device='cpu')\n",
    "y0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b29f3f3-e831-467f-ba59-0dfe0eef6efd",
   "metadata": {},
   "source": [
    "Next, let's consider the setting where you want to pass additional arguments into the forward function because the model is multi-input. Remember that our model can optionally take in `alpha` and `beta` parameters. All we have to do is pass in a tuple of `args` to the `predict` function where each element in `args` is a tensor containing values for one of the inputs to the model.\n",
    "\n",
    "Let's start off by looking at just passing in `alpha` to the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8f0075a9-d112-4d09-af24-5502b5523004",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.2256,  1.3785,  1.2227],\n",
       "        [-0.3800,  0.2527, -0.3178],\n",
       "        [-1.8699, -2.4616, -2.3273],\n",
       "        [ 0.7355,  0.4344,  0.2591],\n",
       "        [-1.1472, -1.0757, -0.7374]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "alpha = torch.randn(5, 1)\n",
    "\n",
    "y + alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "731941ec-9995-4ddd-b7d7-09c8ee108428",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.2256,  1.3785,  1.2227],\n",
       "        [-0.3800,  0.2527, -0.3178],\n",
       "        [-1.8699, -2.4616, -2.3273],\n",
       "        [ 0.7355,  0.4344,  0.2591],\n",
       "        [-1.1472, -1.0757, -0.7374]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(model, X, args=(alpha,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ce8cd6e-4573-4e78-b7e9-4fa249e9c07d",
   "metadata": {},
   "source": [
    "Now, let's try passing in both `alpha` and `beta`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8ff3f489-60e9-4b8d-93f8-3a2ed9e435fa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.3324,  1.4336,  1.3305],\n",
       "        [-0.3165, -0.1477, -0.2999],\n",
       "        [-2.1597, -2.1962, -2.1879],\n",
       "        [ 0.6723,  0.4851,  0.3762],\n",
       "        [-1.0562, -1.0885, -1.2414]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1)\n",
    "beta = torch.randn(5, 1)\n",
    "\n",
    "y * beta + alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "673df291-750a-427d-a738-482bafcc11ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.3324,  1.4336,  1.3305],\n",
       "        [-0.3165, -0.1477, -0.2999],\n",
       "        [-2.1597, -2.1962, -2.1879],\n",
       "        [ 0.6723,  0.4851,  0.3762],\n",
       "        [-1.0562, -1.0885, -1.2414]])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(model, X, args=(alpha, beta))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24750f61-7822-4358-9b06-db8ff8fd51e2",
   "metadata": {},
   "source": [
    "This implementation is extremely flexible. It makes no assumptions on the shape of the underlying data (except that the `batch_size` dimension is the same), and so we could pass in bigger tensors if we wanted to without having to modify the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0c755471-e6eb-4586-bab4-3a25e4206992",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.2256, -0.4559, -2.4970],\n",
       "        [ 0.4818, -0.5384, -1.4230],\n",
       "        [ 0.7123,  0.5552, -0.8678],\n",
       "        [-0.2362, -0.7307, -0.1273],\n",
       "        [-0.9193,  1.1094, -0.7241]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "alpha = torch.randn(5, 3)\n",
    "\n",
    "y + alpha"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ce6aca49-1f1f-4e0a-88d1-1aad00fa66d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.2256, -0.4559, -2.4970],\n",
       "        [ 0.4818, -0.5384, -1.4230],\n",
       "        [ 0.7123,  0.5552, -0.8678],\n",
       "        [-0.2362, -0.7307, -0.1273],\n",
       "        [-0.9193,  1.1094, -0.7241]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict(model, X, args=(alpha,))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cba0b88-72f3-45e6-a37a-d3f5e310a052",
   "metadata": {},
   "source": [
    "This means that if you have a model with one input that is biological sequence and another input that is something more complicated -- like an image, for instance -- you can easily pass both into the model. They just need to be passed in in the same order as defined by the forward function."
   ]
  }
 ],
 "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",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}