Tutorials/Route_Separation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "de8543df-0924-45ef-b4df-eb514d0aa31f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from DirectMultiStep.Models.TensorGen import BeamSearchOptimized as BeamSearch\n",
    "from DirectMultiStep.Models.Configure import prepare_model\n",
    "from DirectMultiStep.Models.Architecture import VanillaTransformerConfig\n",
    "from DirectMultiStep.Utils.Dataset import RoutesStepsSMDataset\n",
    "from DirectMultiStep.Utils.PostProcess import (\n",
    "    BeamResultType, find_valid_paths, process_paths,\n",
    "    find_matching_paths, find_top_n_accuracy, load_pharma_compounds\n",
    ")\n",
    "import torch\n",
    "import yaml\n",
    "from pathlib import Path\n",
    "from rdkit import RDLogger, Chem\n",
    "from tqdm import tqdm\n",
    "\n",
    "RDLogger.DisableLog(\"rdApp.*\")\n",
    "\n",
    "data_path = Path.cwd() / \"Data\"\n",
    "processed_path = data_path / \"Processed\"\n",
    "ckpt_path = data_path / \"Checkpoints\"\n",
    "fig_path = data_path / \"Figures\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d66a95e-5f29-489c-9297-7595f6eb1bb9",
   "metadata": {},
   "source": [
    "# Load Model and Dataset Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f39d3f5e-d41b-4bf5-a908-7a64df4ec497",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with open(processed_path / \"character_dictionary.yaml\", \"rb\") as file:\n",
    "    data = yaml.safe_load(file)\n",
    "    idx_to_token = data[\"invdict\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "249bd768-ead8-4d0b-b1b9-e261f9187a4f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 9,857,333 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "ckpt_name = \"van_6x3_6x3_final.ckpt\"\n",
    "\n",
    "van_enc_conf = VanillaTransformerConfig(\n",
    "    input_dim=53,\n",
    "    output_dim=53,\n",
    "    input_max_length=145 + 135,\n",
    "    output_max_length=1074 + 1,  # 1074 is max length\n",
    "    pad_index=52,\n",
    "    n_layers=6,\n",
    "    ff_mult=3,\n",
    "    attn_bias=False,\n",
    "    ff_activation=\"gelu\",\n",
    "    hid_dim=256,\n",
    ")\n",
    "van_dec_conf = VanillaTransformerConfig(\n",
    "    input_dim=53,\n",
    "    output_dim=53,\n",
    "    input_max_length=145 + 135,\n",
    "    output_max_length=1074 + 1,  # 1074 is max length\n",
    "    pad_index=52,\n",
    "    n_layers=6,\n",
    "    ff_mult=3,\n",
    "    attn_bias=False,\n",
    "    ff_activation=\"gelu\",\n",
    "    hid_dim=256,\n",
    ")\n",
    "model = prepare_model(enc_config=van_enc_conf, dec_config=van_dec_conf)\n",
    "\n",
    "ckpt_torch = torch.load(ckpt_path / ckpt_name, map_location=device)\n",
    "model.load_state_dict(ckpt_torch)\n",
    "model.to(device)\n",
    "model.eval();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa98b9cc-98b4-47ff-a5b5-c76d2a545768",
   "metadata": {},
   "source": [
    "# Load Pharma Compounds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e1184a8a-9670-4636-b6a0-d141087b8745",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "BSObject = BeamSearch(\n",
    "    model=model,\n",
    "    beam_size=50,\n",
    "    start_idx=0,\n",
    "    pad_idx=52,\n",
    "    end_idx=22,\n",
    "    max_length=1074,\n",
    "    idx_to_token=idx_to_token,\n",
    "    device=device,\n",
    ")\n",
    "\n",
    "_products, _sms, _path_strings, _steps_list, nameToIdx = load_pharma_compounds(data_path / \"route_separation.json\")\n",
    "\n",
    "pharma_ds = RoutesStepsSMDataset(\n",
    "        products=_products,\n",
    "        starting_materials=_sms,\n",
    "        path_strings=_path_strings,\n",
    "        n_steps_list=_steps_list,\n",
    "        metadata_path=processed_path / \"character_dictionary.yaml\",\n",
    "    )\n",
    "\n",
    "pharma_dl = torch.utils.data.DataLoader(\n",
    "            dataset=pharma_ds, batch_size=1, shuffle=False, num_workers=0\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a67d3c4-e437-4010-9aef-e531b26975e7",
   "metadata": {},
   "source": [
    "# Generate Routes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b0961dd9-1c13-4c3a-bd83-8a54915a4979",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 36%|███▌      | 386/1072 [00:07<00:13, 52.77it/s]\n",
      " 36%|███▌      | 384/1072 [00:06<00:11, 59.87it/s] \n",
      " 36%|███▌      | 386/1072 [00:06<00:11, 60.06it/s] \n",
      " 34%|███▎      | 360/1072 [00:05<00:11, 62.47it/s] \n",
      " 40%|███▉      | 426/1072 [00:07<00:11, 55.63it/s] \n",
      " 37%|███▋      | 399/1072 [00:06<00:11, 58.06it/s] \n",
      " 39%|███▉      | 419/1072 [00:07<00:11, 55.54it/s] \n"
     ]
    }
   ],
   "source": [
    "all_beam_results_NS2: BeamResultType = []\n",
    "for prod_sm, _, steps in pharma_dl:\n",
    "    beam_result_BS2 = BSObject.decode(\n",
    "        src_BC=prod_sm.to(device), steps_B1=steps.to(device)\n",
    "    )\n",
    "    for beam_result_S2 in beam_result_BS2:\n",
    "        all_beam_results_NS2.append(beam_result_S2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "77f90ddb-4105-47ca-8042-0d78a275a075",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting to canonicalize paths:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 11.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed to canonicalize counter=0 path strings\n",
      "Starting to remove repetitions within beam results:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7/7 [00:00<00:00, 17.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting to find paths with correct product and reactants:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "7it [00:00, 48.58it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Top N Accuracy ---\n",
      "based on raw output:\n",
      "W/o perms: {'Top 1': '0.0', 'Top 3': '0.0', 'Top 5': '0.0', 'Top 10': '0.0', 'Top 20': '0.0', 'Top 50': '0.0'}\n",
      "W/  perms: {'Top 1': '28.6', 'Top 3': '28.6', 'Top 5': '28.6', 'Top 10': '28.6', 'Top 20': '42.9', 'Top 50': '42.9'}\n"
     ]
    }
   ],
   "source": [
    "top_n_vals = [1, 3, 5, 10, 20, 50]\n",
    "\n",
    "valid_paths_NS2n = find_valid_paths(all_beam_results_NS2)\n",
    "correct_paths_NS2n = process_paths(\n",
    "    paths_NS2n=valid_paths_NS2n, true_products=_products, true_reacs=_sms, commercial_stock=None, verbose=True\n",
    "        )\n",
    "match_accuracy_N, perm_match_accuracy_N = find_matching_paths(\n",
    "    correct_paths_NS2n, _path_strings\n",
    ")\n",
    "\n",
    "freqs_noperm = find_top_n_accuracy(match_accuracy_N, top_n_vals)\n",
    "freqs_wperm = find_top_n_accuracy(perm_match_accuracy_N, top_n_vals)\n",
    "print(\"---- Top N Accuracy ---\")\n",
    "print(\"based on raw output:\")\n",
    "print(f\"W/o perms: {freqs_noperm}\")\n",
    "print(f\"W/  perms: {freqs_wperm}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "109e7c45-1154-4dfc-b0eb-3cf23e3d8c3e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ranks: [None, 1, None, None, 1, None, 18]\n",
      "Route names and their corresponding indices (depends on # of SMs):\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'First Half': [0, 1, 2], 'Second Half': [3, 4, 5, 6]}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(\"Ranks:\", perm_match_accuracy_N)\n",
    "print(\"Route names and their corresponding indices (depends on # of SMs):\")\n",
    "nameToIdx"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "antonEnv",
   "language": "python",
   "name": "env"
  },
  "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": 5
}