test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2b6dc8e0-0478-45eb-be59-f0b85a0fa1f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/shared/morgan/miniforge3/envs/acegen/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from copy import deepcopy\n",
    "from typing import Callable, Union\n",
    "from functools import partial\n",
    "\n",
    "import torch\n",
    "from tensordict.nn import TensorDictModule\n",
    "from tensordict.nn.probabilistic import set_interaction_type as set_exploration_type\n",
    "from tensordict.tensordict import TensorDictBase\n",
    "from torchrl.collectors import RandomPolicy\n",
    "from torchrl.envs import EnvBase, TensorDictPrimer, TransformedEnv, CatFrames, InitTracker, StepCounter, UnsqueezeTransform\n",
    "from torchrl.envs.utils import ExplorationType, step_mdp\n",
    "\n",
    "\n",
    "from acegen.rl_env.smiles_env import SMILESEnv\n",
    "from acegen.vocabulary import SMILESVocabulary, SMILESTokenizer, SMILESTokenizer2\n",
    "from acegen.data.utils import smiles_to_tensordict\n",
    "from acegen.models import create_gru_actor, adapt_state_dict, create_gpt2_actor\n",
    "\n",
    "from promptsmiles import ScaffoldDecorator, FragmentLinker\n",
    "_has_promptsmiles = True\n",
    "\n",
    "from acegen.rl_env.utils import generate_complete_smiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "82a57db8-1c41-4dff-9b81-25f3bdbf8473",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true,
     "source_hidden": true
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/shared/morgan/miniforge3/envs/acegen/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'acegen.vocabulary.tokenizers.SMILESTokenizer'>\n",
      "['C', 'C', '1', '=', 'C', '(', 'C', '(', '=', 'O', ')', 'N', '2', 'C', 'C', 'C', 'C', 'C', '2', '=', 'N', '1', ')', 'C', 'C', 'N', '3', 'C', 'C', 'C', '(', 'C', 'C', '3', ')', 'C', '4', '=', 'N', 'O', 'C', '5', '=', 'C', '4', 'C', '=', 'C', 'C', '(', '=', 'C', '5', ')', 'F']\n",
      "CC1=C(C(=O)N2CCCCC2=N1)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F\n",
      "<class 'acegen.vocabulary.tokenizers.SMILESTokenizer2'>\n",
      "['C', 'C', '1', '=', 'C', '(', 'C', '(', '=', 'O', ')', 'N', '2', 'C', 'C', 'C', 'C', 'C', '2', '=', 'N', '1', ')', 'C', 'C', 'N', '3', 'C', 'C', 'C', '(', 'C', 'C', '3', ')', 'C', '4', '=', 'N', 'O', 'C', '5', '=', 'C', '4', 'C', '=', 'C', 'C', '(', '=', 'C', '5', ')', 'F']\n",
      "CC1=C(C(=O)N2CCCCC2=N1)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F\n",
      "<class 'acegen.vocabulary.tokenizers.DeepSMILESTokenizer'>\n",
      "['C', 'C', '=', 'C', 'C', '=', 'O', ')', 'N', 'C', 'C', 'C', 'C', 'C', '6', '=', 'N', '%', '1', '0', ')', ')', ')', ')', ')', ')', ')', ')', 'C', 'C', 'N', 'C', 'C', 'C', 'C', 'C', '6', ')', ')', 'C', '=', 'N', 'O', 'C', '=', 'C', '5', 'C', '=', 'C', 'C', '=', 'C', '6', ')', 'F']\n",
      "CC2=C(C(=O)N1CCCCC1=N2)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F\n",
      "<class 'acegen.vocabulary.tokenizers.AISTokenizer'>\n",
      "['[CH3;!R;C]', '[c;R;CCN]', '1', '[n;R;CC]', '[c;R;CNN]', '2', '[n;R;CCC]', '(', '[c;R;CNO]', '(', '=', '[O;!R;C]', ')', '[c;R;CCC]', '1', '[CH2;!R;CC]', '[CH2;!R;CN]', '[N;R;CCC]', '1', '[CH2;R;CN]', '[CH2;R;CC]', '[CH;R;CCC]', '(', '[c;R;CCN]', '3', '[n;R;CO]', '[o;R;CN]', '[c;R;CCO]', '4', '[cH;R;CC]', '[c;R;CCF]', '(', '[F;!R;C]', ')', '[cH;R;CC]', '[cH;R;CC]', '[c;R;CCC]', '3', '4', ')', '[CH2;R;CC]', '[CH2;R;CN]', '1', ')', '[CH2;R;CN]', '[CH2;R;CC]', '[CH2;R;CC]', '[CH2;R;CC]', '2']\n",
      "Cc1nc2n(c(=O)c1CCN1CCC(c3noc4cc(F)ccc34)CC1)CCCC2\n",
      "<class 'acegen.vocabulary.tokenizers.SAFETokenizer'>\n",
      "['C', 'c', '1', 'n', 'c', '2', 'n', '(', 'c', '(', '=', 'O', ')', 'c', '1', '7', ')', 'C', 'C', 'C', 'C', '2', '.', 'c', '1', '8', 'n', 'o', 'c', '2', 'c', 'c', '(', 'F', ')', 'c', 'c', 'c', '1', '2', '.', 'N', '1', '6', 'C', 'C', 'C', '8', 'C', 'C', '1', '.', 'C', '7', 'C', '6']\n",
      "Cc1nc2n(c(=O)c1CCN1CCC(c3noc4cc(F)ccc34)CC1)CCCC2\n",
      "<class 'acegen.vocabulary.tokenizers.SELFIESTokenizer'>\n",
      "['[C]', '[C]', '[=C]', '[Branch1]', '[S]', '[C]', '[=Branch1]', '[C]', '[=O]', '[N]', '[C]', '[C]', '[C]', '[C]', '[C]', '[Ring1]', '[=Branch1]', '[=N]', '[Ring1]', '[O]', '[C]', '[C]', '[N]', '[C]', '[C]', '[C]', '[Branch1]', '[Branch1]', '[C]', '[C]', '[Ring1]', '[=Branch1]', '[C]', '[=N]', '[O]', '[C]', '[=C]', '[Ring1]', '[Branch1]', '[C]', '[=C]', '[C]', '[=Branch1]', '[Ring2]', '[=C]', '[Ring1]', '[=Branch1]', '[F]']\n",
      "CC1=C(C(=O)N2CCCCC2=N1)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F\n"
     ]
    }
   ],
   "source": [
    "from acegen.vocabulary import tokenizers\n",
    "\n",
    "for tokenizer in [tokenizers.SMILESTokenizer, tokenizers.SMILESTokenizer2, tokenizers.DeepSMILESTokenizer, tokenizers.AISTokenizer, tokenizers.SAFETokenizer, tokenizers.SELFIESTokenizer]:\n",
    "    print(tokenizer)\n",
    "    smiles = \"CC1=C(C(=O)N2CCCCC2=N1)CCN3CCC(CC3)C4=NOC5=C4C=CC(=C5)F\"\n",
    "    TOK = tokenizer()\n",
    "    tokens = TOK.tokenize(smiles)\n",
    "    print(tokens)\n",
    "    untokens = TOK.untokenize(tokens)\n",
    "    print(untokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "addf09e6-62cd-4d8d-8012-007ec952d5d8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/shared/morgan/acegen-open/acegen/rl_env/utils.py\u001b[0m(268)\u001b[0;36mgenerate_complete_smiles\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    266 \u001b[0;31m        \u001b[0;32mif\u001b[0m \u001b[0mprompt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    267 \u001b[0;31m            \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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[0;32m--> 268 \u001b[0;31m            \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprompt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstr\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[0;32m    269 \u001b[0;31m                \u001b[0mprompt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mprompt\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    270 \u001b[0;31m            \u001b[0;31m# Encode the prompt(s)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ipdb>  c\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/shared/morgan/acegen-open/acegen/rl_env/utils.py\u001b[0m(355)\u001b[0;36mgenerate_complete_smiles\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    353 \u001b[0;31m            \u001b[0mtensordicts\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"next\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"done\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m~\u001b[0m\u001b[0mfinished\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\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[0;32m    354 \u001b[0;31m        \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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[0;32m--> 355 \u001b[0;31m        \u001b[0moutput_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensordicts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontiguous\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[0;32m    356 \u001b[0;31m        \u001b[0moutput_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrefine_names\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"time\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    357 \u001b[0;31m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ipdb>  c\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/shared/morgan/acegen-open/acegen/rl_env/utils.py\u001b[0m(360)\u001b[0;36mgenerate_complete_smiles\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    358 \u001b[0;31m    \u001b[0;32mif\u001b[0m \u001b[0mreturn_smiles_only\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    359 \u001b[0;31m        \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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[0;32m--> 360 \u001b[0;31m        \u001b[0msmiles\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutput_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"action\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\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[0;32m    361 \u001b[0;31m        \u001b[0msmiles_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mvocabulary\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msmi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0msmi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msmiles\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    362 \u001b[0;31m        \u001b[0;31m# Replace failed encodings with original prompt for PromptSMILES\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ipdb>  c\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/shared/morgan/acegen-open/acegen/rl_env/utils.py\u001b[0m(268)\u001b[0;36mgenerate_complete_smiles\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    266 \u001b[0;31m        \u001b[0;32mif\u001b[0m \u001b[0mprompt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    267 \u001b[0;31m            \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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[0;32m--> 268 \u001b[0;31m            \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprompt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstr\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[0;32m    269 \u001b[0;31m                \u001b[0mprompt\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mprompt\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mbatch_size\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    270 \u001b[0;31m            \u001b[0;31m# Encode the prompt(s)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ipdb>  c\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/shared/morgan/acegen-open/acegen/rl_env/utils.py\u001b[0m(355)\u001b[0;36mgenerate_complete_smiles\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    353 \u001b[0;31m            \u001b[0mtensordicts\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"next\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"done\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m~\u001b[0m\u001b[0mfinished\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\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[0;32m    354 \u001b[0;31m        \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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[0;32m--> 355 \u001b[0;31m        \u001b[0moutput_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensordicts\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontiguous\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[0;32m    356 \u001b[0;31m        \u001b[0moutput_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrefine_names\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m...\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"time\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    357 \u001b[0;31m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ipdb>  c\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/shared/morgan/acegen-open/acegen/rl_env/utils.py\u001b[0m(360)\u001b[0;36mgenerate_complete_smiles\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    358 \u001b[0;31m    \u001b[0;32mif\u001b[0m \u001b[0mreturn_smiles_only\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    359 \u001b[0;31m        \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\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[0;32m--> 360 \u001b[0;31m        \u001b[0msmiles\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutput_data\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"action\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\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[0;32m    361 \u001b[0;31m        \u001b[0msmiles_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mvocabulary\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdecode\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msmi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0msmi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0msmiles\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    362 \u001b[0;31m        \u001b[0;31m# Replace failed encodings with original prompt for PromptSMILES\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n"
     ]
    },
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "ipdb>  c\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorDict(\n",
      "    fields={\n",
      "        action: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.float32, is_shared=False),\n",
      "        done: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "        is_init: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "        mask: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "        next: TensorDict(\n",
      "            fields={\n",
      "                done: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "                is_init: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "                observation: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.int32, is_shared=False),\n",
      "                reward: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.float32, is_shared=False),\n",
      "                sequence: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.int32, is_shared=False),\n",
      "                terminated: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "                truncated: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False)},\n",
      "            batch_size=torch.Size([2, 200]),\n",
      "            device=cpu,\n",
      "            is_shared=False),\n",
      "        observation: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.int32, is_shared=False),\n",
      "        promptsmiles: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.float32, is_shared=False),\n",
      "        sample_log_prob: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.float32, is_shared=False),\n",
      "        sequence: Tensor(shape=torch.Size([2, 200]), device=cpu, dtype=torch.int32, is_shared=False),\n",
      "        terminated: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False),\n",
      "        truncated: Tensor(shape=torch.Size([2, 200, 1]), device=cpu, dtype=torch.bool, is_shared=False)},\n",
      "    batch_size=torch.Size([2, 200]),\n",
      "    device=cpu,\n",
      "    is_shared=False)\n",
      "NC(=N)c1cc2nc(-c3ccc(Cn4cc(COc5ccc(-c6ccc(OCc7cn(Cc8ccc(-c9nc%10cc(C(=N)N)ccc%10[nH]9)o8)nn7)cc6)cc5)nn4)o3)[nH]c2cc1\n",
      "c1cc(OCc2cn(Cc3ccc(-c4nc5cc(C(=N)N)ccc5[nH]4)o3)nn2)ccc1-c6ccc(OCc7cn(Cc8ccc(-c9nc%10cc(C(=N)N)ccc%10[nH]9)o8)nn7)cc6)CC1=O\n"
     ]
    }
   ],
   "source": [
    "output_data = generate_complete_smiles(\n",
    "    environment=env,\n",
    "    vocabulary=vocabulary,\n",
    "    policy_sample=actor_inference,\n",
    "    policy_evaluate=actor_training,\n",
    "    prompt=None, #\"c%11ccccc%11\", # \"c1ccccc1C(=O)\",\n",
    "    promptsmiles=\"N=C(N)c1c(*)cc2[nH]c(-c3ccc(Cn4cc(COc5ccc(-c6c(*)cc(OCc7cn(Cc8ccc(-c9nc%10cc(C(=N)N)ccc%10[nH]9)o8)nn7)cc6)cc5)nn4)o3)nc2c1\", #\"c%12ccc(*)cc%12\", # \"N1(*)CCN(CC1)CCCCN(*)\", # fragment \"N1(*)CCNCC1.C1(*)CC1.c1cncc(*)c1\", # scaffold \"N1(*)CCN(CC1)CCCCN(*)\"\n",
    "    promptsmiles_optimize=False,\n",
    "    promptsmiles_shuffle=True,\n",
    "    promptsmiles_multi=False,\n",
    "    return_smiles_only=False\n",
    ")\n",
    "if isinstance(output_data, list):\n",
    "    _ = [print(s) for s in output_data]\n",
    "else:\n",
    "    print(output_data)\n",
    "    smiles = output_data.get(\"action\").cpu()\n",
    "    smiles_str = [vocabulary.decode(smi.numpy()) for smi in smiles]\n",
    "    _ = [print(s) for s in smiles_str]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4e11d1dd-5912-40dc-9201-d81852849783",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Load prereqs\n",
    "smiles = [\n",
    "    'c1ccccc1',\n",
    "    'c1ccccc1C(=O)C'\n",
    "]\n",
    "\n",
    "# Get available device\n",
    "device = torch.device(\"cpu\")\n",
    "create_actor, voc_path, ckpt_path, tokenizer = (create_gru_actor, \"priors/chembl_filtered_vocabulary.txt\", \"priors/gru_chembl_filtered.ckpt\", SMILESTokenizer())\n",
    "#create_actor, voc_path, ckpt_path, tokenizer = (create_gpt2_actor, \"priors/enamine_real_vocabulary.txt\",\"/shared/albert/acegen-open/priors/gpt2_enamine_real.ckpt\", SMILESTokenizer2())\n",
    "# Load vocabulary\n",
    "with open(voc_path, \"r\") as f:\n",
    "        tokens = f.read().splitlines()\n",
    "tokens_dict = dict(zip(tokens, range(len(tokens))))\n",
    "vocabulary = SMILESVocabulary.create_from_dict(tokens_dict, start_token=\"GO\", end_token=\"EOS\", tokenizer=tokenizer)\n",
    "# Create models\n",
    "ckpt = torch.load(ckpt_path)\n",
    "actor_training, actor_inference = create_actor(vocabulary_size=len(vocabulary))\n",
    "actor_inference.load_state_dict(\n",
    "    adapt_state_dict(ckpt, actor_inference.state_dict())\n",
    ")\n",
    "actor_training.load_state_dict(adapt_state_dict(ckpt, actor_training.state_dict()))\n",
    "actor_inference = actor_inference.to(device)\n",
    "actor_training = actor_training.to(device)\n",
    "prior = deepcopy(actor_training)\n",
    " # Create RL environment\n",
    "# For RNNs, create a transform to populate initial tensordict with recurrent states equal to 0.0\n",
    "rhs_primers = []\n",
    "if hasattr(actor_training, \"rnn_spec\"):\n",
    "    primers = actor_training.rnn_spec.expand(2)\n",
    "    rhs_primers.append(TensorDictPrimer(primers))\n",
    "env_kwargs = {\n",
    "    \"start_token\": vocabulary.start_token_index,\n",
    "    \"end_token\": vocabulary.end_token_index,\n",
    "    \"length_vocabulary\": len(vocabulary),\n",
    "    \"batch_size\": 2,\n",
    "    \"max_length\": 200,\n",
    "    \"device\": device,\n",
    "}\n",
    "def create_env_fn():\n",
    "    \"\"\"Create a single RL rl_env.\"\"\"\n",
    "    env = SMILESEnv(**env_kwargs)\n",
    "    env = TransformedEnv(env)\n",
    "    env.append_transform(StepCounter())\n",
    "    env.append_transform(InitTracker())\n",
    "    for rhs_primer in rhs_primers:\n",
    "        env.append_transform(rhs_primer)\n",
    "    return env\n",
    "\n",
    "env = create_env_fn()\n",
    "#data = generate_complete_smiles(policy=actor_inference, environment=env)"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}