{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prober example\n",
    "\n",
    "This example shows how to use `Prober` to couple already trained probes to a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch # PyTorch\n",
    "\n",
    "from pytorch_probing import Prober, ParallelModuleDict # Prober and dictionary of modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start creating a example model, a simple MLP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExampleModel(torch.nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super().__init__()\n",
    "\n",
    "        self.linear1 = torch.nn.Linear(input_size, hidden_size)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.linear2 = torch.nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.linear2(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ExampleModel(\n",
       "  (linear1): Linear(in_features=2, out_features=3, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (linear2): Linear(in_features=3, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_size = 2\n",
    "hidden_size = 3\n",
    "output_size = 1\n",
    "\n",
    "model = ExampleModel(input_size, hidden_size, output_size)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And a probe. Any torch module can be used as a probe. In this example we gonna use a simple Linear layer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=3, out_features=2, bias=True)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probe_size = 2\n",
    "\n",
    "probe = torch.nn.Linear(hidden_size, probe_size)\n",
    "probe.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we created the Prober, passing it the model and a dictionary mapping the paths of the modules to the probes that must be coupled to its outputs. When a `None` value is passed, it creates a `Identity` module, that just pass its inputs to the outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "Create Prober"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prober(\n",
       "  (_module): ExampleModel(\n",
       "    (linear1): InterceptorLayer(\n",
       "      (_module): Linear(in_features=2, out_features=3, bias=True)\n",
       "    )\n",
       "    (relu): InterceptorLayer(\n",
       "      (_module): ReLU()\n",
       "    )\n",
       "    (linear2): Linear(in_features=3, out_features=1, bias=True)\n",
       "  )\n",
       "  (_probes): ModuleDict(\n",
       "    (linear1): Linear(in_features=3, out_features=2, bias=True)\n",
       "    (relu): Identity()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probes = {\"linear1\":probe, \"relu\":None}\n",
    "\n",
    "probed_model = Prober(model, probes)\n",
    "probed_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observe that the \"Prober\" modifies the original model in-place:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [
     "View modified model"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ExampleModel(\n",
       "  (linear1): InterceptorLayer(\n",
       "    (_module): Linear(in_features=2, out_features=3, bias=True)\n",
       "  )\n",
       "  (relu): InterceptorLayer(\n",
       "    (_module): ReLU()\n",
       "  )\n",
       "  (linear2): Linear(in_features=3, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We pass a sample value to the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = torch.randn([10, 2])\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs = probed_model(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the output is a tuple with the model output in the first value, and the probes outputs in the second:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.3165],\n",
       "        [-0.3262],\n",
       "        [-0.3362],\n",
       "        [-0.4985],\n",
       "        [-0.2987],\n",
       "        [-0.3520],\n",
       "        [-0.3182],\n",
       "        [-0.3269],\n",
       "        [-0.3418],\n",
       "        [-0.3352]])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": [
     "View probes outputs"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'linear1': tensor([[ 0.1408,  0.8607],\n",
       "         [ 0.1683,  0.9637],\n",
       "         [-0.7003,  0.6100],\n",
       "         [ 0.6062,  1.2991],\n",
       "         [-0.1432,  0.7140],\n",
       "         [-1.0194,  0.2385],\n",
       "         [-0.4221,  0.6943],\n",
       "         [-0.6234,  0.5710],\n",
       "         [-0.9021,  0.4412],\n",
       "         [ 0.1946,  1.0828]]),\n",
       " 'relu': tensor([[0.0000, 0.0346, 0.0000],\n",
       "         [0.0000, 0.0534, 0.0000],\n",
       "         [1.3431, 0.0000, 0.0000],\n",
       "         [0.0000, 0.3868, 0.0000],\n",
       "         [0.0000, 0.0000, 0.0000],\n",
       "         [1.3982, 0.0000, 0.3337],\n",
       "         [0.7008, 0.0000, 0.0000],\n",
       "         [1.0109, 0.0000, 0.0000],\n",
       "         [1.5421, 0.0000, 0.0000],\n",
       "         [0.0000, 0.0708, 0.0000]])}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple probes in the same place\n",
    "\n",
    "We can also use more than one probe in the same place. For showing it, we gonna create a second probe and reduce the probed model to the original model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Linear(in_features=3, out_features=1, bias=True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "probe2_size = 1\n",
    "\n",
    "probe2 = torch.nn.Linear(hidden_size, probe2_size)\n",
    "probe2.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = probed_model.reduce()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can than create a `ParallelModuleDict` with the two probes. We called with some input, the `ParallelModuleDict` pass the input to all its modules, and return a dictionary with each module output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear1_probes = ParallelModuleDict({\"probe1\":probe, \"probe2\":probe2})\n",
    "\n",
    "probes = {\"linear1\":linear1_probes}\n",
    "\n",
    "probed_model = Prober(model, probes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = torch.randn([10, 2])\n",
    "\n",
    "with torch.no_grad():\n",
    "    outputs2 = probed_model(inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "View probes outputs 2"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'linear1': {'probe1': tensor([[-0.2682,  0.8801],\n",
       "          [ 0.0351,  0.7865],\n",
       "          [-1.2988,  0.1559],\n",
       "          [-0.3947,  0.7546],\n",
       "          [-0.3776,  0.7478],\n",
       "          [-0.4901,  0.4542],\n",
       "          [ 0.1075,  0.9543],\n",
       "          [-0.5112,  0.7165],\n",
       "          [-0.5990,  0.6198],\n",
       "          [-0.1853,  0.7252]]),\n",
       "  'probe2': tensor([[-0.4589],\n",
       "          [-0.3497],\n",
       "          [-0.3711],\n",
       "          [-0.4320],\n",
       "          [-0.4252],\n",
       "          [-0.3211],\n",
       "          [-0.4073],\n",
       "          [-0.4413],\n",
       "          [-0.4184],\n",
       "          [-0.3721]])}}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs2[1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}