LlaMol_SMILES_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load LlaMol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "import rdkit\n",
    "from rdkit import Chem\n",
    "import rdkit.rdBase as rkrb\n",
    "import rdkit.RDLogger as rkl\n",
    "import os\n",
    "import torch \n",
    "import logging\n",
    "import numpy as np\n",
    "from plot_utils import check_metrics\n",
    "from sample import Sampler\n",
    "import pandas as pd\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "if \"cuda\" in device:\n",
    "    # dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'\n",
    "    dtype = \"float16\" if torch.cuda.is_available() else \"float32\"\n",
    "else:\n",
    "    dtype = \"float32\"\n",
    "\n",
    "logger = rkl.logger()\n",
    "logger.setLevel(rkl.ERROR)\n",
    "rkrb.DisableLog(\"rdApp.error\")\n",
    "\n",
    "torch.set_num_threads(8)\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# sampler\n",
    "sampler = Sampler(\n",
    "    load_path=os.path.join(\n",
    "        os.getcwd(), \"out\", \"llama2-M-Full-RSS-Canonical.pt\"\n",
    "    ),\n",
    "    device=device,\n",
    "    seed=1234,\n",
    "    dtype=dtype,\n",
    "    compile=False,\n",
    ")\n",
    "\n",
    "# setting   \n",
    "num_samples = 100\n",
    "df_comp = pd.read_parquet(os.path.join(os.getcwd(),\"data\",\"OrganiX13.parquet\"))\n",
    "df_comp = df_comp.sample(n=2_500_000)\n",
    "comp_context_dict = {c: df_comp[c].to_numpy() for c in [\"logp\", \"sascore\", \"mol_weight\"]} \n",
    "comp_smiles = df_comp[\"smiles\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert SMILES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "from typing import List, Dict\n",
    "import json\n",
    "from rdkit.Chem import AllChem\n",
    "\n",
    "@torch.no_grad()\n",
    "def convert_to_chemiscope(smiles_list : List[str], context_dict : Dict[str, List[float]]):\n",
    "    # For more details on the file format: https://chemiscope.org/docs/tutorial/input-reference.html\n",
    "\n",
    "    structures = []\n",
    "    remove_list = []\n",
    "    for i,smi in enumerate(smiles_list):\n",
    "        mol = Chem.MolFromSmiles(smi)\n",
    "        if mol is None:\n",
    "            logging.info(f\"Mol invalid: {smi} ! Skipping...\")\n",
    "            remove_list.append(i)\n",
    "            continue\n",
    "\n",
    "        res = AllChem.EmbedMolecule(mol,randomSeed=0xf00d, maxAttempts=20)\n",
    "        # res = AllChem.Compute2DCoords(mol)\n",
    "\n",
    "        if res != 0:\n",
    "            logging.info(f\"Could not calculate coordinates for {smi}! Skipping..\")\n",
    "            remove_list.append(i)\n",
    "            continue\n",
    "        \n",
    "\n",
    "        conf = list(mol.GetConformers())[0]\n",
    "        x,y,z = [],[],[]\n",
    "        symbols = []\n",
    "        for atom, coords in zip(mol.GetAtoms(), conf.GetPositions()):\n",
    "            symbols.append(atom.GetSymbol())\n",
    "            x.append(coords[0])\n",
    "            y.append(coords[1])\n",
    "            z.append(coords[2])\n",
    "        \n",
    "        structures.append({\n",
    "            \"size\": len(x),\n",
    "            \"names\": symbols,\n",
    "            \"x\": x,\n",
    "            \"y\": y,\n",
    "            \"z\" : z\n",
    "        })\n",
    "\n",
    "    properties = {}\n",
    "    \n",
    "    for c in context_dict:\n",
    "        properties[c] = {\n",
    "            \"target\": \"structure\",\n",
    "            \"values\": [v for i, v in enumerate(context_dict[c]) if i not in remove_list]\n",
    "        }\n",
    "        \n",
    "        \n",
    "    data = {\n",
    "        \"meta\": {\n",
    "            # // the name of the dataset\n",
    "            \"name\": \"Test Dataset\",\n",
    "            # // description of the dataset, OPTIONAL\n",
    "            \"description\": \"This contains data from generated molecules\",\n",
    "            # // authors of the dataset, OPTIONAL\n",
    "            \"authors\": [\"Niklas Dobberstein, niklas.dobberstein@scai.fraunhofer.de\"],\n",
    "            # // references for the dataset, OPTIONAL\n",
    "            \"references\": [\n",
    "                \"\",\n",
    "            ],\n",
    "        \n",
    "        },\n",
    "        \"properties\": properties,\n",
    "        \"structures\": structures\n",
    "    }\n",
    "    \n",
    "    out_path = os.path.join(os.getcwd(), \"chemiscope_gen.json\")\n",
    "    with open(out_path, \"w\") as f:\n",
    "        json.dump(data, f)\n",
    "\n",
    "    logging.info(f\"Wrote file {out_path}\")\n",
    "\n",
    "convert_to_chemiscope([\n",
    "    \"CC=O\",\n",
    "    \"s1ccnc1\"\n",
    "], {\"logp\": [1.0,2.0], \"sascore\": [1.5,-2.0]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Condition: SMILES token sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "context_smi = \"C=C(C)C(=O)O[11CH3]\"  # Define your SMILES input if needed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import Draw\n",
    "import os\n",
    "\n",
    "from plot_utils import calc_context_from_smiles\n",
    "\n",
    "# Set parameters\n",
    "context_cols_options = [\"logp\", \"sascore\", \"mol_weight\"]\n",
    "\n",
    "\n",
    "temperature = 0.8\n",
    "num_samples = 40  # Define number of samples\n",
    "device = 'cpu'  # Define device for PyTorch\n",
    "\n",
    "# Define sliders and parameters for logp, sascore, and mol_weight\n",
    "logp = 0.0\n",
    "sascore = 2.0\n",
    "mol_weight = 3.0\n",
    "\n",
    "# Define context dictionary\n",
    "selected_context_cols = [\"logp\", \"sascore\", \"mol_weight\"]\n",
    "context_dict = {}\n",
    "for c in selected_context_cols:\n",
    "    if c == \"logp\":\n",
    "        val = logp\n",
    "    elif c == \"sascore\":\n",
    "        val = sascore\n",
    "    else:\n",
    "        val = mol_weight\n",
    "    val = round(val, 2)\n",
    "    context_dict[c] = val * torch.ones((num_samples,), device=device, dtype=torch.float)\n",
    "\n",
    "# Generate SMILES using the provided context\n",
    "smiles, context = sampler.generate(\n",
    "    context_cols=context_dict,\n",
    "    context_smi=context_smi,\n",
    "    start_smiles=None,\n",
    "    num_samples=num_samples,\n",
    "    max_new_tokens=256,\n",
    "    temperature=temperature,\n",
    "    top_k=25,\n",
    "    total_gen_steps=int(np.ceil(num_samples / 1000)),\n",
    "    return_context=True\n",
    ")\n",
    "\n",
    "# Save SMILES to gen_smiles.txt\n",
    "with open(\"gen_smiles.txt\", \"w\") as f:\n",
    "    for s in smiles:\n",
    "        f.write(f\"{s}\\n\")\n",
    "\n",
    "# Display SMILES as RDKit molecules\n",
    "def display_molecules(smiles_list, context_dict):\n",
    "    molecules = [Chem.MolFromSmiles(smiles) for smiles in smiles_list]\n",
    "    \n",
    "    # Convert RDKit molecules to images and store them in a list\n",
    "    images = [Draw.MolToImage(mol) for mol in molecules]\n",
    "\n",
    "    # Create a subplot grid to display the images\n",
    "    num_images = len(images)\n",
    "    num_cols = 1  # Number of columns in the grid\n",
    "    num_rows = (num_images + num_cols - 1) // num_cols  # Calculate the number of rows\n",
    "\n",
    "    fig, axes = plt.subplots(num_rows, num_cols, figsize=(150, 150))\n",
    "    fig.subplots_adjust(hspace=0.5)\n",
    "\n",
    "    calculated_context = {c: [] for c in context_dict}\n",
    "    for i, ax in enumerate(axes.flat):\n",
    "        if i < num_images:\n",
    "            ax.imshow(images[i])\n",
    "            for j, c in enumerate(context_dict):\n",
    "                smiles = smiles_list[i]\n",
    "                smi_con = round(calc_context_from_smiles([smiles], c)[0], 2)\n",
    "                calculated_context[c].append(smi_con)\n",
    "                ax.text(0.5, -0.1 * j, f\"{c}: {context_dict[c][i]} vs {smi_con}\", \n",
    "                        transform=ax.transAxes, fontsize=10, ha='center')\n",
    "\n",
    "            ax.axis('off')\n",
    "        else:\n",
    "            fig.delaxes(ax)  # Remove empty subplots if there are more rows than images\n",
    "\n",
    "    plt.savefig(\"gen_mols.png\")\n",
    "    plt.close()\n",
    "\n",
    "# Generate molecule display\n",
    "display_molecules(smiles, context_dict)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}