BeamSearch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle \n",
    "import yaml\n",
    "import pandas as pd\n",
    "from PrepareData import prepare_data\n",
    "\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim, Tensor\n",
    "from torch.nn import functional as F\n",
    "import pickle \n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import rdkit\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import AllChem\n",
    "from rdkit.Chem import Draw\n",
    "import seaborn as sns\n",
    "from architecture import CLIP\n",
    "from tqdm import tqdm\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "dtype = torch.float32\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from train_utils import load_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_deterministic(random_seed = 0):\n",
    "    torch.manual_seed(random_seed)\n",
    "    torch.cuda.manual_seed(random_seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    np.random.seed(random_seed)\n",
    "    \n",
    "make_deterministic(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best_epoch 424\n",
      "best_clip_epoch 286\n",
      "best_recon_epoch 589\n",
      "best_total_loss 0.09221624111450172\n",
      "best_clip_loss 0.012083162560423872\n",
      "best_recon_loss 0.07850571382272092\n"
     ]
    }
   ],
   "source": [
    "config = yaml.safe_load(open('./checkpoints/TEST_RUN_REPRODUCE/config.yaml', 'r'))\n",
    "logs = pickle.load(open('./checkpoints/TEST_RUN_REPRODUCE/logs.pickle', 'rb'))\n",
    "for key in logs:\n",
    "    if \"best\" in key:\n",
    "        print(key, logs[key])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# config['data']['batch_size'] = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57257af2d3744a689be0fd7330981a09",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65785dc69f0d42f19e28e3accf919557",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SMILES WILL BE RANDOMIZED\n"
     ]
    }
   ],
   "source": [
    "model = load_model(config['train']['checkpoint_dir'], type=\"best_latest\")\n",
    "model.eval()\n",
    "dataloaders, max_charge, num_species, scaler = prepare_data(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_ids = pickle.load(open('./checkpoints/TEST_RUN_REPRODUCE/val_ids.pickle', 'rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "79it [00:00, 95.19it/s] \n"
     ]
    }
   ],
   "source": [
    "all_ids = []\n",
    "with torch.no_grad():\n",
    "    for i, data in tqdm(enumerate(dataloaders['val'])):    \n",
    "        data = {k: v.to(device) for k, v in data.items()}\n",
    "        all_ids.append(data['index'].detach().cpu())\n",
    "all_ids = torch.cat(all_ids, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert((all_ids.sort()[0] == val_ids.sort()[0]).sum() == all_ids.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): CLIP(\n",
       "    (Molecule_Encoder): EGNN(\n",
       "      (embedding): Linear(in_features=15, out_features=256, bias=True)\n",
       "      (embedding_out): Linear(in_features=256, out_features=15, bias=True)\n",
       "      (gcl_0): E_GCL_mask(\n",
       "        (edge_mlp): Sequential(\n",
       "          (0): Linear(in_features=513, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "          (3): SiLU()\n",
       "        )\n",
       "        (node_mlp): Sequential(\n",
       "          (0): Linear(in_features=527, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (att_mlp): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=1, bias=True)\n",
       "          (1): Sigmoid()\n",
       "        )\n",
       "        (act_fn): SiLU()\n",
       "      )\n",
       "      (gcl_1): E_GCL_mask(\n",
       "        (edge_mlp): Sequential(\n",
       "          (0): Linear(in_features=513, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "          (3): SiLU()\n",
       "        )\n",
       "        (node_mlp): Sequential(\n",
       "          (0): Linear(in_features=527, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (att_mlp): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=1, bias=True)\n",
       "          (1): Sigmoid()\n",
       "        )\n",
       "        (act_fn): SiLU()\n",
       "      )\n",
       "      (gcl_2): E_GCL_mask(\n",
       "        (edge_mlp): Sequential(\n",
       "          (0): Linear(in_features=513, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "          (3): SiLU()\n",
       "        )\n",
       "        (node_mlp): Sequential(\n",
       "          (0): Linear(in_features=527, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (att_mlp): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=1, bias=True)\n",
       "          (1): Sigmoid()\n",
       "        )\n",
       "        (act_fn): SiLU()\n",
       "      )\n",
       "      (gcl_3): E_GCL_mask(\n",
       "        (edge_mlp): Sequential(\n",
       "          (0): Linear(in_features=513, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "          (3): SiLU()\n",
       "        )\n",
       "        (node_mlp): Sequential(\n",
       "          (0): Linear(in_features=527, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (att_mlp): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=1, bias=True)\n",
       "          (1): Sigmoid()\n",
       "        )\n",
       "        (act_fn): SiLU()\n",
       "      )\n",
       "      (gcl_4): E_GCL_mask(\n",
       "        (edge_mlp): Sequential(\n",
       "          (0): Linear(in_features=513, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "          (3): SiLU()\n",
       "        )\n",
       "        (node_mlp): Sequential(\n",
       "          (0): Linear(in_features=527, out_features=256, bias=True)\n",
       "          (1): SiLU()\n",
       "          (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "        )\n",
       "        (att_mlp): Sequential(\n",
       "          (0): Linear(in_features=256, out_features=1, bias=True)\n",
       "          (1): Sigmoid()\n",
       "        )\n",
       "        (act_fn): SiLU()\n",
       "      )\n",
       "      (node_dec): Sequential(\n",
       "        (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "        (1): SiLU()\n",
       "        (2): Linear(in_features=256, out_features=256, bias=True)\n",
       "      )\n",
       "      (graph_dec): Sequential(\n",
       "        (0): Linear(in_features=256, out_features=256, bias=True)\n",
       "        (1): SiLU()\n",
       "        (2): Linear(in_features=256, out_features=512, bias=True)\n",
       "      )\n",
       "    )\n",
       "    (Spectra_Encoder): ViT(\n",
       "      (proc): Sequential(\n",
       "        (0): Unfold(kernel_size=(1, 7), dilation=1, padding=0, stride=(1, 7))\n",
       "        (1): Transpose()\n",
       "        (2): Linear(in_features=7, out_features=512, bias=True)\n",
       "      )\n",
       "      (enc): ViTransformerEncoder(\n",
       "        (layers): ModuleList(\n",
       "          (0): ViTransformerEncoderLayer(\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (mha): MultiHeadAttention(\n",
       "              (heads): ModuleList(\n",
       "                (0): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (1): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (2): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (3): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (4): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (5): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (6): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "              )\n",
       "              (linear): Linear(in_features=511, out_features=512, bias=True)\n",
       "            )\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "              (1): GELU(approximate='none')\n",
       "              (2): Dropout(p=0.1, inplace=False)\n",
       "              (3): Linear(in_features=1024, out_features=512, bias=True)\n",
       "            )\n",
       "          )\n",
       "          (1): ViTransformerEncoderLayer(\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (mha): MultiHeadAttention(\n",
       "              (heads): ModuleList(\n",
       "                (0): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (1): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (2): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (3): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (4): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (5): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (6): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "              )\n",
       "              (linear): Linear(in_features=511, out_features=512, bias=True)\n",
       "            )\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "              (1): GELU(approximate='none')\n",
       "              (2): Dropout(p=0.1, inplace=False)\n",
       "              (3): Linear(in_features=1024, out_features=512, bias=True)\n",
       "            )\n",
       "          )\n",
       "          (2): ViTransformerEncoderLayer(\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (mha): MultiHeadAttention(\n",
       "              (heads): ModuleList(\n",
       "                (0): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (1): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (2): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (3): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (4): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (5): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (6): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "              )\n",
       "              (linear): Linear(in_features=511, out_features=512, bias=True)\n",
       "            )\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "              (1): GELU(approximate='none')\n",
       "              (2): Dropout(p=0.1, inplace=False)\n",
       "              (3): Linear(in_features=1024, out_features=512, bias=True)\n",
       "            )\n",
       "          )\n",
       "          (3): ViTransformerEncoderLayer(\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (mha): MultiHeadAttention(\n",
       "              (heads): ModuleList(\n",
       "                (0): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (1): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (2): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (3): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (4): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (5): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (6): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "              )\n",
       "              (linear): Linear(in_features=511, out_features=512, bias=True)\n",
       "            )\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "              (1): GELU(approximate='none')\n",
       "              (2): Dropout(p=0.1, inplace=False)\n",
       "              (3): Linear(in_features=1024, out_features=512, bias=True)\n",
       "            )\n",
       "          )\n",
       "          (4): ViTransformerEncoderLayer(\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (mha): MultiHeadAttention(\n",
       "              (heads): ModuleList(\n",
       "                (0): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (1): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (2): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (3): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (4): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (5): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "                (6): Head(\n",
       "                  (q_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (k_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                  (v_lin): Linear(in_features=512, out_features=73, bias=False)\n",
       "                )\n",
       "              )\n",
       "              (linear): Linear(in_features=511, out_features=512, bias=True)\n",
       "            )\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (ffn): Sequential(\n",
       "              (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "              (1): GELU(approximate='none')\n",
       "              (2): Dropout(p=0.1, inplace=False)\n",
       "              (3): Linear(in_features=1024, out_features=512, bias=True)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (pos_emb): Embedding(1000, 512)\n",
       "        (dropout_emb): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (mlp): Linear(in_features=512, out_features=512, bias=True)\n",
       "    )\n",
       "    (smiles_decoder): LatentToMol(\n",
       "      (embed): Embedding(25, 512)\n",
       "      (pe): PositionalEncodings(\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (classifier): Linear(in_features=512, out_features=25, bias=True)\n",
       "      (trfmencoder): TransformerEncoder(\n",
       "        (layers): ModuleList(\n",
       "          (0): TransformerEncoderLayer(\n",
       "            (self_attn): MultiheadAttention(\n",
       "              (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "            (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout1): Dropout(p=0.1, inplace=False)\n",
       "            (dropout2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (1): TransformerEncoderLayer(\n",
       "            (self_attn): MultiheadAttention(\n",
       "              (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "            (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout1): Dropout(p=0.1, inplace=False)\n",
       "            (dropout2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (2): TransformerEncoderLayer(\n",
       "            (self_attn): MultiheadAttention(\n",
       "              (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "            )\n",
       "            (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "            (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
       "            (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "            (dropout1): Dropout(p=0.1, inplace=False)\n",
       "            (dropout2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "      (specencoder): LatentEncoder(\n",
       "        (featurizer): Sequential(\n",
       "          (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "          (1): GELU(approximate='none')\n",
       "          (2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "          (3): GELU(approximate='none')\n",
       "        )\n",
       "      )\n",
       "      (ln_f): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "      (drop): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = load_model(config['train']['checkpoint_dir'], type=\"best_total\")\n",
    "model = model.eval()\n",
    "model.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data = next(iter(dataloaders['val']))\n",
    "# data = {k: v.to(device) for k, v in data.items()}\n",
    "# mol_latents, spec_latents, smile_preds, logit_scale, ids = model(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from train_utils import calculate_decoder_accuracy\n",
    "# acc = calculate_decoder_accuracy(model, dataloaders, k=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Faster Sampling algorithm for greedy and random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "PAD = 0\n",
    "UNK = 1\n",
    "EOS = 2\n",
    "SOS = 3\n",
    "MASK = 4\n",
    "\n",
    "SMI_MAX_SIZE = 71\n",
    "\n",
    "class Sampler():\n",
    "    def __init__(self, model, vocab):\n",
    "        self.model = model\n",
    "        self.vocab = vocab\n",
    "        self.max_len = 70\n",
    "    def sample(self, embed, greedy_decode=False):\n",
    "        samples = []\n",
    "        # print(embed.shape)\n",
    "        with torch.no_grad():\n",
    "            embed = embed.to(device)\n",
    "            smiles_seq = torch.full((embed.shape[0], 1), SOS).long().to(device)\n",
    "            \n",
    "            for i in range(SMI_MAX_SIZE):\n",
    "                logits = self.model.forward(embed, smiles_seq)\n",
    "                # print(\"logits\", logits.shape)\n",
    "                probs = F.softmax(logits[:,-1], dim= -1)\n",
    "                # print(probs.shape)\n",
    "                # break\n",
    "                if greedy_decode:\n",
    "                    pred_id = torch.argmax(probs, dim= -1)\n",
    "                    pred_id = pred_id.unsqueeze(1)\n",
    "                else:\n",
    "                    pred_id = torch.multinomial(probs, num_samples=1)\n",
    "                # print(pred_id.shape)\n",
    "                # break\n",
    "                smiles_seq = torch.cat([smiles_seq, pred_id], dim=1)\n",
    "                \n",
    "            for i in range(len(smiles_seq)):\n",
    "                smile = self.vocab.from_seq(smiles_seq[i].cpu().numpy())\n",
    "                final_smile = \"\"\n",
    "                for char in smile[1:]: # first is start token\n",
    "                    if char == \"<eos>\" :\n",
    "                        break\n",
    "                    final_smile += char\n",
    "                try:\n",
    "                    final_smile = Chem.CanonSmiles(final_smile)\n",
    "                except:\n",
    "                    pass\n",
    "                samples.append(final_smile)\n",
    "                \n",
    "        return samples\n",
    "  \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_decoder_accuracy( model, dataloaders, k=1, greedy_decode=True):\n",
    "    pred_smiles_list = []\n",
    "    og_smiles_list = []\n",
    "    validity = []\n",
    "    validity_count = 0\n",
    "    count = 0\n",
    "    sampler = Sampler(model.module.smiles_decoder, model.module.vocab)\n",
    "    all_samples = []\n",
    "    og_samples = []\n",
    "    with torch.no_grad():\n",
    "        for i, data in tqdm(enumerate(dataloaders['val'])):\n",
    "            data = {k: v.to(device) for k, v in data.items()}\n",
    "            spec_latents = model.module.forward_spec(data)\n",
    "            generated_samples = sampler.sample(spec_latents, greedy_decode=greedy_decode)\n",
    "            all_samples += generated_samples\n",
    "            \n",
    "            for og in data['smiles']:\n",
    "                og_smile = \"\"\n",
    "                chars = model.module.vocab.from_seq(og)\n",
    "                for char in chars:\n",
    "                    if char != \"<pad>\" and char != \"<eos>\" and char != \"<sos>\" and char != \"<unk>\":\n",
    "                        og_smile += char\n",
    "                try:\n",
    "                    og_smile = Chem.CanonSmiles(og_smile)\n",
    "                except:\n",
    "                    og_smile=None\n",
    "                og_samples.append(og_smile)\n",
    "          \n",
    "            print(\"No of Hits : \", (np.array(og_samples) == np.array(all_samples)).sum() / len(og_samples))\n",
    "\n",
    "        return count / len(og_smiles_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# acc = calculate_decoder_accuracy(model, dataloaders, k=1, greedy_decode=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Beam Search Sampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "PAD = 0\n",
    "UNK = 1\n",
    "EOS = 2\n",
    "SOS = 3\n",
    "MASK = 4\n",
    "\n",
    "class BeamSampler():\n",
    "    def __init__(self, model, vocab, beam_size=5):\n",
    "        self.model = model\n",
    "        self.vocab = vocab\n",
    "        self.max_len = 70\n",
    "        self.beam_size = beam_size\n",
    "\n",
    "    def sample(self, embed, greedy_decode=False):\n",
    "        embed = embed.unsqueeze(0).to(device) # shape (1, 512) \n",
    "        self.model.eval()\n",
    "        sample_tensor = torch.zeros((1,self.max_len), dtype=torch.int64).to(device)\n",
    "        sample_tensor[0,0] = SOS\n",
    "        beams = [(sample_tensor, 0.0)]\n",
    "\n",
    "        for i in range(self.max_len -1):\n",
    "            \n",
    "            new_beams = []\n",
    "            for beam in beams:\n",
    "                \n",
    "                vec, score = beam[0], beam[1]\n",
    "                tensor = vec[:,:i+1]\n",
    "                logits = self.model.forward(embed, tensor)[:,-1,:]\n",
    "                probabilities = F.log_softmax(logits, dim=1)\n",
    "                topk_values, topk_indices = torch.topk(probabilities, self.beam_size, dim=-1)\n",
    "                for value, ind in zip(topk_values[0], topk_indices[0]):\n",
    "                    new_vec = vec.clone()\n",
    "                    new_vec[0,i+1] = ind\n",
    "                    new_score = score + value\n",
    "                    new_beams.append((new_vec, new_score))     \n",
    "     \n",
    "            beams = sorted(new_beams, key=lambda x:x[1].item(), reverse=True)[:self.beam_size]\n",
    "\n",
    "        sampled_smiles = [] \n",
    "        for beam in beams:\n",
    "            smiles = \"\"\n",
    "            chars = self.vocab.from_seq(beam[0][0])\n",
    "            for char in chars:\n",
    "                if char != \"<pad>\" and char != \"<eos>\" and char != \"<sos>\" and char != \"<unk>\":\n",
    "                    smiles += char\n",
    "            sampled_smiles.append(smiles)\n",
    "                \n",
    "        return sampled_smiles\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sampler = BeamSampler(model.module.smiles_decoder, model.module.vocab)\n",
    "# sample_smiles = sampler.sample(spec_latents[0], greedy_decode=False)\n",
    "# print(\"sample_smile\", sample_smile)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_decoder_accuracy( model, dataloaders, beam_size=5):\n",
    "    with torch.no_grad():\n",
    "        pred_smiles_list = []\n",
    "        og_smiles_list = []\n",
    "        count = 0\n",
    "        sampler = BeamSampler(model.module.smiles_decoder, model.module.vocab, beam_size=beam_size)\n",
    "        \n",
    "        for i, data in tqdm(enumerate(dataloaders['val'])):\n",
    "            data = {k: v.to(device) for k, v in data.items()}\n",
    "            spec_latents = model.module.forward_spec(data)\n",
    "            for spec, og in zip(spec_latents, data['smiles'] ):\n",
    "                ls = [sampler.sample(spec)[0]] # checking only the top beam\n",
    "                # ls = sampler.sample(spec) # checking all beams\n",
    "                generated_smiles = []\n",
    "                for smi in ls:\n",
    "                    try:\n",
    "                        generated_smiles.append(Chem.CanonSmiles(smi))\n",
    "                    except:\n",
    "                        pass\n",
    "                og_smile = \"\"\n",
    "                chars = model.module.vocab.from_seq(og)\n",
    "                for char in chars:\n",
    "                    if char != \"<pad>\" and char != \"<eos>\" and char != \"<sos>\" and char != \"<unk>\":\n",
    "                        og_smile += char\n",
    "                try:\n",
    "                    og_smile = Chem.CanonSmiles(og_smile)\n",
    "                except:\n",
    "                    og_smile=None\n",
    "                    \n",
    "                if og_smile is not None and og_smile in generated_smiles:\n",
    "                    count += 1\n",
    "                \n",
    "                og_smiles_list.append(og_smile)\n",
    "                pred_smiles_list.append(generated_smiles)\n",
    "            print(\"No of Hits : \",count / len(og_smiles_list))\n",
    "            # if sampling takes too long, we can stop after sampling around 5000 molecules    \n",
    "            if i == 5:\n",
    "                break\n",
    "        \n",
    "        return count / len(og_smiles_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [00:28, 28.95s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.4765625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [00:55, 27.46s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5234375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "3it [01:21, 26.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.51953125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4it [01:46, 26.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5205078125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [02:10, 25.24s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5171875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [02:37, 31.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5149739583333334\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "acc = calculate_decoder_accuracy(model, dataloaders, beam_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [01:14, 74.81s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5859375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [02:28, 74.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.580078125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "3it [03:41, 73.73s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5911458333333334\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4it [04:54, 73.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5771484375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [06:07, 73.40s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5640625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [07:20, 88.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5670572916666666\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "acc = calculate_decoder_accuracy(model, dataloaders, beam_size=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [02:05, 125.78s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.56640625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [04:17, 129.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.55859375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "3it [06:22, 127.40s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5677083333333334\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4it [08:28, 126.80s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.576171875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [10:33, 125.94s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.58125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [12:37, 151.42s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.5859375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "acc = calculate_decoder_accuracy(model, dataloaders, beam_size=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1it [04:45, 285.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.49609375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2it [09:18, 278.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.54296875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "3it [13:58, 278.85s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.55078125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4it [18:33, 277.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.546875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [23:09, 276.80s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.55\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "5it [27:47, 333.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No of Hits :  0.560546875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "acc = calculate_decoder_accuracy(model, dataloaders, beam_size=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sbdd-env",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}