Upload X-LoRA-Gemma_Inference.ipynb

X-LoRA-Gemma_Inference.ipynb

@@ -0,0 +1,1101 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "3288987d",
+   "metadata": {},
+   "source": [
+    "# X-LoRA Inference: Gemma-7b model for molecular design \n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "25beb240-1ae1-4537-9cc6-da621862d0bd",
+   "metadata": {},
+   "source": [
+    "### Helper functions "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2c18b20-b1a9-4f3e-ae84-2a551e2ed69c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import random\n",
+    "\n",
+    "import torch\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "import transformers\n",
+    "from datasets import load_dataset\n",
+    "from datasets import IterableDataset\n",
+    "\n",
+    "from transformers import Trainer\n",
+    "from transformers import TrainingArguments\n",
+    "from transformers import DataCollatorWithPadding\n",
+    "from transformers import TrainerCallback\n",
+    "from transformers import AutoConfig\n",
+    "from transformers import BitsAndBytesConfig\n",
+    "\n",
+    "from peft import LoraConfig, get_peft_model\n",
+    "from torch.utils.data import Dataset\n",
+    "from transformers import get_linear_schedule_with_warmup\n",
+    "from accelerate import infer_auto_device_map\n",
+    "import math\n",
+    "import numpy as np\n",
+    "import unidecode\n",
+    "import pandas as pd\n",
+    "from matplotlib import pyplot as plt\n",
+    "import peft\n",
+    "\n",
+    "from tqdm.notebook import tqdm\n",
+    "\n",
+    "device='cuda'\n",
+    "\n",
+    "def params(model):\n",
+    "    model_parameters = filter(lambda p: p.requires_grad, model.parameters())\n",
+    "    params = sum([np.prod(p.size()) for p in model_parameters])\n",
+    "\n",
+    "    print(\"Number of model arameters: \", params) \n",
+    "\n",
+    "def generate_response (model,tokenizer,text_input=\"Biology offers amazing\",\n",
+    "                      num_return_sequences=1,\n",
+    "                      temperature=1., #the higher the temperature, the more creative the model becomes\n",
+    "                      max_new_tokens=127,\n",
+    "                      num_beams=1,\n",
+    "                      top_k = 50,\n",
+    "                      top_p =0.9,repetition_penalty=1.,eos_token_id=107,verbatim=False,\n",
+    "                      exponential_decay_length_penalty_fac=None,add_special_tokens  =True, eos_token=None, \n",
+    "                      ):\n",
+    "\n",
+    "    if eos_token==None:\n",
+    "        eos_token=tokenizer('<end_of_turn>', add_special_tokens  =False,   ) ['input_ids'][0]\n",
+    "        \n",
+    "    inputs = tokenizer(text_input,  \n",
+    "                              add_special_tokens  =add_special_tokens,  \n",
+    "                              return_tensors ='pt').to(device)\n",
+    "    if verbatim:\n",
+    "        print (\"Length of input, tokenized: \", inputs[\"input_ids\"].shape, inputs[\"input_ids\"],\"eos_token: \", eos_token)\n",
+    "    with torch.no_grad():\n",
+    "          outputs = model.generate(#input_ids=inputs.to(device), \n",
+    "                                   input_ids = inputs[\"input_ids\"],\n",
+    "                                    attention_mask = inputs[\"attention_mask\"] , # This is usually done automatically by the tokenizer\n",
+    "                                    max_new_tokens=max_new_tokens,\n",
+    "                                   temperature=temperature, #value used to modulate the next token probabilities.\n",
+    "                                   num_beams=num_beams,\n",
+    "                                   top_k = top_k,\n",
+    "                                   top_p = top_p,\n",
+    "                                   num_return_sequences = num_return_sequences,\n",
+    "                                   eos_token_id=eos_token,\n",
+    "                                   pad_token_id = eos_token,\n",
+    "                                   do_sample =True, \n",
+    "                                   repetition_penalty=repetition_penalty, \n",
+    "                                  )\n",
+    "\n",
+    "    return tokenizer.batch_decode(outputs[:,inputs[\"input_ids\"].shape[1]:].detach().cpu().numpy(), skip_special_tokens=True)\n",
+    "\n",
+    "def generate_answer (model,tokenizer,system='You a helpful assistant. You are familiar with materials science. ',\n",
+    "                     q='What is spider silk in the context of bioinspired materials?',\n",
+    "                        repetition_penalty=1.1,\n",
+    "                                           top_p=0.1, top_k=32,  \n",
+    "                                  temperature=.6,max_new_tokens=512, verbatim=False, eos_token=None,add_special_tokens=True,\n",
+    "                     prepend_response='', messages=[],\n",
+    "                                ):\n",
+    "\n",
+    "    if eos_token==None:\n",
+    "        eos_token= tokenizer.eos_token_id\n",
+    "        \n",
+    "    if system==None:\n",
+    "        messages.append ({\"role\": \"user\", \"content\": q} )\n",
+    "    else:\n",
+    "        messages.append ({\"role\": \"user\", \"content\": system+q})\n",
+    "         \n",
+    "    txt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, )\n",
+    "    txt=txt+prepend_response\n",
+    "     \n",
+    "    output_text=generate_response (model,tokenizer,text_input=txt,eos_token_id=eos_token,\n",
+    "                                  num_return_sequences=1,  repetition_penalty=repetition_penalty,\n",
+    "                                           top_p=top_p, top_k=top_k,  add_special_tokens  =add_special_tokens,\n",
+    "                                \n",
+    "                                 temperature=temperature,max_new_tokens=max_new_tokens, verbatim=verbatim, \n",
+    "                                            \n",
+    "                                           )\n",
+    "    return (  output_text[0] )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "75d89d27-8386-4859-a36e-ce4842415b59",
+   "metadata": {},
+   "source": [
+    "### Load X-LoRA Gemma model "
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "cd1b66f6-1fe1-4b2c-9309-fe01d34d7d54",
+   "metadata": {},
+   "source": [
+    "https://github.com/EricLBuehler/xlora"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12848c38-cc0c-41c7-bf04-9856730458df",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from xlora.xlora_utils import load_model  \n",
+    "\n",
+    "XLoRa_model_name = 'lamm-mit/x-lora-gemma-7b'\n",
+    "\n",
+    "model, tokenizer=load_model(model_name = XLoRa_model_name, \n",
+    "                           device='cuda:0',\n",
+    "                           use_flash_attention_2=True, \n",
+    "                           dtype=torch.bfloat16,\n",
+    "                            )\n",
+    "eos_token_id= tokenizer('<end_of_turn>', add_special_tokens=False, ) ['input_ids'][0]\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b197ffd5-7752-4081-9227-c46a485afeec",
+   "metadata": {},
+   "source": [
+    "### Inference using Guidance "
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "f7009898-17a9-468a-970a-59d7c80553ca",
+   "metadata": {},
+   "source": [
+    "https://github.com/guidance-ai/guidance"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "80b62bf2-a424-4858-a321-f55e3327b070",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from guidance import models\n",
+    "from guidance import gen, select, system, user, assistant,  newline\n",
+    "from IPython.display import display, Markdown\n",
+    "\n",
+    "gpt = models.TransformersChat(model=model, tokenizer=tokenizer)\n",
+    "gpt_question_asker = gpt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1cb5a867-a127-45c2-b75b-35883a78930b",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "with user():        \n",
+    "    lm =gpt + f\"\"\"List the most important biomolecules used in biological materials to make polymers with multifunctional qualities.\"\"\" \n",
+    "\n",
+    "with assistant():        \n",
+    "    lm+=\"[\"+gen('res1', max_tokens=1024)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a841c58c-bded-4741-80df-66ca434bfac0",
+   "metadata": {},
+   "source": [
+    "### Inference using Hugging Face generate functions  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26a27dc2-4e28-4fee-b37c-446281cd23da",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "system_prompt='You are an expert in biological molecular engineering. '\n",
+    "q=\"\"\"\n",
+    "What are potential molecular engineering approaches to create better materials? Name specific molecules of interest.\n",
+    "\"\"\"\n",
+    "\n",
+    "res=generate_answer (model, tokenizer,system=system_prompt,\n",
+    "                     q=q,\n",
+    "                     repetition_penalty=1.,  top_p=0.9, top_k=256,  \n",
+    "                     temperature=.5,max_new_tokens=512, verbatim=False, \n",
+    "                )\n",
+    "\n",
+    "display (Markdown (\"## X-LoRA:\\n\\n\"+res))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "82d162fe-6149-44d4-afbe-63213b10f183",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "system_prompt='You are an expert in biological molecular engineering. '\n",
+    "q=\"\"\"\n",
+    "List the most important biomolecules used in biological materials to make polymers with multifunctional qualities.\n",
+    "\"\"\"\n",
+    "messages=[]\n",
+    "res=generate_answer (model, tokenizer,system=system_prompt,\n",
+    "                     q=q,   repetition_penalty=1., top_p=0.9, top_k=256,    temperature=.5,max_new_tokens=512, verbatim=False,messages=messages  )\n",
+    "\n",
+    "display (Markdown (\"## X-LoRA:\\n\\n\"+res))\n",
+    "messages.append ({\"role\": \"assistant\", \"content\": res} )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f9ad8c01-7cfd-4017-a8af-0b72b4ea25fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "system_prompt=None\n",
+    "q=\"\"\"\n",
+    "How does chitin form a material, specifically in terms of molecular interactions? \n",
+    "\"\"\" \n",
+    "res=generate_answer (model, tokenizer,system=system_prompt,\n",
+    "                     q=q,   repetition_penalty=1., top_p=0.9, top_k=256,    temperature=.1,max_new_tokens=512, verbatim=False,messages=messages,\n",
+    "                )\n",
+    "\n",
+    "display (Markdown (\"## X-LoRA:\\n\\n\"+res))\n",
+    "messages.append ({\"role\": \"assistant\", \"content\": res} )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6f520fd9-0d06-4971-9b58-74d8d2c3e2ef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "system_prompt=None\n",
+    "q=\"\"\"\n",
+    "Thank you. What are potential chemical modifications of N-acetylglucosamine units that would improve mechanical properties?\n",
+    "\"\"\" \n",
+    "res=generate_answer (model, tokenizer,system=system_prompt,\n",
+    "                     q=q,   repetition_penalty=1., top_p=0.9, top_k=256,    temperature=.1,max_new_tokens=512, verbatim=False,messages=messages,\n",
+    "                )\n",
+    "\n",
+    "display (Markdown (\"## X-LoRA:\\n\\n\"+res))\n",
+    "messages.append ({\"role\": \"assistant\", \"content\": res} )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ce5a2293-b66d-4ef6-987e-451dc1a92621",
+   "metadata": {},
+   "source": [
+    "### Molecule design examples"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e547bed4-da94-48c7-b9dd-00da7732ef20",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "import pandas as pd\n",
+    "from sklearn.preprocessing import MinMaxScaler\n",
+    "\n",
+    "df_smiles=pd.read_csv ('./QM9.csv')\n",
+    "SMILES_LIST=list (df_smiles['smiles'])\n",
+    "\n",
+    "X = df_smiles.iloc[:, 0].values.reshape(-1, 1)  # Input feature, reshaped for compatibility\n",
+    "y = df_smiles.iloc[:, 1:]  # Target features\n",
+    "\n",
+    "# Scaling the target features\n",
+    "scaler = MinMaxScaler()\n",
+    "y_scaled = scaler.fit_transform(y)\n",
+    "\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "\n",
+    "X_train, X_test, y_train, y_test= train_test_split(X, y_scaled, test_size=0.2, random_state=42)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "44c43109-0606-42d2-a1b6-01278ff6432f",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "import seaborn as sns\n",
+    "from sklearn.metrics import mean_squared_error\n",
+    "labels = [\"mu\", \"alpha\", \"homo\", \"lumo\", \"gap\", \"r2\", \"zpve\", \"cv\", \"u0\", \"u298\", \"h298\", \"g298\"]\n",
+    "\n",
+    "def return_str(vals=np.array ([.1, .5, .6, 2.])):\n",
+    "    ch=''\n",
+    "    for i in range (len (vals)):\n",
+    "        ch=ch+f'{vals[i]:1.3f},'\n",
+    "        \n",
+    "    return ch[:-1]   \n",
+    "\n",
+    "def extract_start_and_end(string_input, start_token='[', end_token=']'):\n",
+    "    \"\"\"\n",
+    "    Extracts the substring from 'string_input' that is enclosed between the first occurrence of\n",
+    "    'start_token' and the last occurrence of 'end_token'.\n",
+    "\n",
+    "    Args:\n",
+    "    string_input (str): The string from which to extract the substring.\n",
+    "    start_token (str): The starting delimiter. Default is '['.\n",
+    "    end_token (str): The ending delimiter. Default is ']'.\n",
+    "\n",
+    "    Returns:\n",
+    "    str: The extracted substring. If 'start_token' or 'end_token' is not found, returns an empty string.\n",
+    "    \"\"\"\n",
+    "    # Find the index of the first occurrence of start_token\n",
+    "    i = string_input.find(start_token)\n",
+    "    # Find the index of the last occurrence of end_token\n",
+    "    j = string_input.rfind(end_token)\n",
+    "\n",
+    "    # Check if both tokens are found and i < j to ensure proper enclosure\n",
+    "    if i == -1 or j == -1 or i >= j:\n",
+    "        return \"\"\n",
+    "    else:\n",
+    "        # Extract and return the content between the first start_token and the last end_token\n",
+    "        return string_input[i + 1:j]\n",
+    "\n",
+    "def is_SMILES_novel (SMILES, SMILES_LIST=None):\n",
+    "\n",
+    "    if SMILES_LIST !=None:\n",
+    "        \n",
+    "        if SMILES not in SMILES_LIST:\n",
+    "            is_novel=True\n",
+    "        else:\n",
+    "            is_novel=False\n",
+    "    else:\n",
+    "        is_novel=None\n",
+    "    return is_novel\n",
+    "    \n",
+    "def visualize_SMILES (smiles_code, dir_path='./' , root='', sample_count=0):\n",
+    "    molecule = Chem.MolFromSmiles(smiles_code)\n",
+    "                    \n",
+    "    # Generate an image of the molecule\n",
+    "    molecule_image = Draw.MolToImage(molecule)\n",
+    "    \n",
+    "    # Display the image directly in Jupyter Notebook\n",
+    "    display(molecule_image)\n",
+    "    \n",
+    "    image_path=f\"{dir_path}/SMILES_{sample_count}_{root}_molecule_image.png\"\n",
+    "    molecule_image.save(image_path)\n",
+    "\n",
+    "    return image_path\n",
+    "\n",
+    "\n",
+    "def design_from_target(\n",
+    "    model,\n",
+    "    tokenizer,\n",
+    "    target,\n",
+    "    temperature=0.1,\n",
+    "    num_beams=1,\n",
+    "    top_k=50,\n",
+    "    top_p=0.95,\n",
+    "    repetition_penalty=1.0,\n",
+    "    messages=[]\n",
+    "):\n",
+    "    # Format the target line for molecular property generation\n",
+    "    line = f'GenerateMolecularProperties<{return_str(target)}>'\n",
+    "    \n",
+    "    # Add the line to the message history\n",
+    "    messages.append({\"role\": \"user\", \"content\": line})\n",
+    "    \n",
+    "    # Apply chat template with optional tokenization\n",
+    "    line = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
+    "    \n",
+    "    # Generate response with specified parameters\n",
+    "    result = generate_response(\n",
+    "        model,\n",
+    "        tokenizer,\n",
+    "        text_input=line,\n",
+    "        num_return_sequences=1,\n",
+    "        temperature=temperature,\n",
+    "        top_k=top_k,\n",
+    "        top_p=top_p,\n",
+    "        max_new_tokens=256\n",
+    "    )[0]\n",
+    "    \n",
+    "    return result\n",
+    "\n",
+    "def properties_from_SMILES(\n",
+    "    model,\n",
+    "    tokenizer,\n",
+    "    target,\n",
+    "    temperature=0.1,\n",
+    "    top_k=128,\n",
+    "    top_p=0.9,\n",
+    "    num_beams=1,\n",
+    "    repetition_penalty=1.0\n",
+    "):\n",
+    "    # Format the target line for molecular property calculation\n",
+    "    line = f'CalculateMolecularProperties<{target}>'\n",
+    "    \n",
+    "    # Initialize messages and add the formatted line\n",
+    "    messages = [{\"role\": \"user\", \"content\": line}]\n",
+    "    \n",
+    "    # Apply chat template with optional tokenization\n",
+    "    line = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
+    "    \n",
+    "    # Generate response with specified parameters\n",
+    "    result = generate_response(\n",
+    "        model,\n",
+    "        tokenizer,\n",
+    "        text_input=line,\n",
+    "        num_return_sequences=1,\n",
+    "        temperature=temperature,\n",
+    "        top_k=top_k,\n",
+    "        top_p=top_p,\n",
+    "        max_new_tokens=256\n",
+    "    )[0]\n",
+    "    \n",
+    "    # Extract relevant part of the result and convert to float list\n",
+    "    result = extract_start_and_end(result, start_token='[', end_token=']')\n",
+    "    return [float(i) for i in result.split(',')]\n",
+    "\n",
+    "  \n",
+    "def avg_properties_from_SMILES (model, tokenizer, SMILES ='O=C(N)C1OC(CO)C(O)C(O)C1O', SMILES_dir='./',\n",
+    "                                temperature=0.01, top_k=50,top_p=0.95, num_beams=1,   repetition_penalty=1.,\n",
+    "                                labels=None, N_prop=6, plot_results=True):\n",
+    "    if not os.path.exists(SMILES_dir):\n",
+    "        os.makedirs(SMILES_dir)    \n",
+    "    properties=[]\n",
+    "    if labels==None and plot_results:\n",
+    "        labels= ['mu',\n",
+    "                 'alpha',\n",
+    "                 'homo',\n",
+    "                 'lumo',\n",
+    "                 'gap',\n",
+    "                 'r2',\n",
+    "                 'zpve',\n",
+    "                 'cv',\n",
+    "                 'u0',\n",
+    "                 'u298',\n",
+    "                 'h298',\n",
+    "                 'g298']\n",
+    "    successful=0\n",
+    "    for i in tqdm(range (N_prop)):\n",
+    "        \n",
+    "        try:\n",
+    "            _prop=properties_from_SMILES (model, tokenizer, SMILES,temperature=temperature, top_k=top_k,top_p=top_p,\n",
+    "                                          num_beams=num_beams, repetition_penalty=repetition_penalty,\n",
+    "                                         )\n",
+    "            if len (_prop)==len (labels):\n",
+    "        \n",
+    "                properties.append(np.array( _prop) )\n",
+    "                successful+=1\n",
+    "        except:\n",
+    "            print (end=\"\")\n",
+    "             \n",
+    "    all_properties = np.array(properties)\n",
+    "    \n",
+    "    # Calculate mean and standard deviation for each property\n",
+    "    means = np.mean(all_properties, axis=0)\n",
+    "    std_devs = np.std(all_properties, axis=0)\n",
+    "    \n",
+    "    # Labels for the x-axis\n",
+    "    if plot_results: \n",
+    "        # Creating the plot with error bars\n",
+    "        plt.figure(figsize=(6, 4))\n",
+    "        plt.errorbar(labels, means, yerr=std_devs, fmt='o', ecolor='red', capsize=5, capthick=2, marker='s', color='blue')\n",
+    "        plt.xticks(rotation=45)\n",
+    "        plt.xlabel('Property')\n",
+    "        plt.ylabel('Value')\n",
+    "        plt.title('Average Properties with Error Bars')\n",
+    "        plt.tight_layout()\n",
+    "        plt.savefig(SMILES_dir + f\"avg_prop_{SMILES}.svg\", format=\"svg\")\n",
+    "        \n",
+    "        plt.show()\n",
+    "        print (f\"Successful attempts: {successful}/{N_prop}\")\n",
+    "        \n",
+    "    return means, std_devs        \n",
+    "\n",
+    "def is_valid_smiles(smiles):\n",
+    "    # This function tries to create a molecule object from a SMILES string.\n",
+    "    # If the molecule object is created successfully and is not None, the SMILES is valid.\n",
+    "    mol = Chem.MolFromSmiles(smiles)\n",
+    "    return mol is not None\n",
+    "    \n",
+    "def design_molecule(model, tokenizer, target=None, temperature=0.1,\n",
+    "             num_beams=1,top_k=50,top_p=0.95,  repetition_penalty=1.,\n",
+    "                   SMILES_LIST=None, dir_path='./', messages=[],N_attempts_for_forward=1):\n",
+    "\n",
+    "    if not os.path.exists(dir_path):\n",
+    "        os.makedirs(dir_path)\n",
+    "    if target.any()==None:\n",
+    "        target = np.random.rand(12)\n",
+    "    \n",
+    "    try:\n",
+    "        SMILES=design_from_target (model, tokenizer, target, messages=messages)\n",
+    "    except:\n",
+    "        SMILES=None\n",
+    "        print (\"Generation failed.\")\n",
+    "\n",
+    "    is_novel=is_SMILES_novel (SMILES, SMILES_LIST)\n",
+    "    print (\"Result: \", SMILES, \"is novel: \", is_novel, \"is valid: \", is_valid_smiles(SMILES))\n",
+    "    try:\n",
+    "        visualize_SMILES (SMILES, dir_path=dir_path)\n",
+    "    except:\n",
+    "        print (\"Vis failed.\")\n",
+    "\n",
+    "    try:\n",
+    "        if N_attempts_for_forward==1:\n",
+    "            predicted = properties_from_SMILES(model, tokenizer, SMILES,temperature_pred, num_beams,\n",
+    "                                            top_k, top_p, repetition_penalty)\n",
+    "        else:\n",
+    "            predicted,_=avg_properties_from_SMILES(model, tokenizer, SMILES, SMILES_dir=SMILES_dir,\n",
+    "                                temperature=temperature_pred, top_k=top_k,top_p=top_p, num_beams=num_beams,   repetition_penalty=repetition_penalty,\n",
+    "                                labels=labels, N_prop=N_attempts_for_forward, plot_results=False)\n",
+    "\n",
+    "        sns.set_style(\"whitegrid\")\n",
+    "        plt.gcf().set_facecolor('white')\n",
+    "        # Assuming GT_res and predictions are your data arrays/lists for Ground Truth and Predictions respectively\n",
+    "        \n",
+    "        x = np.arange(len(labels))  # Label locations\n",
+    "        width = 0.35  # Width of the bars\n",
+    "        \n",
+    "        fig, ax = plt.subplots(figsize=(9, 5))\n",
+    "        rects1 = ax.bar(x - width/2, target, width, label='Target')\n",
+    "        rects2 = ax.bar(x + width/2, predicted, width, label='Predicted properties')\n",
+    "        \n",
+    "        # Add some text for labels, title and custom x-axis tick labels, etc.\n",
+    "        ax.set_ylabel('Values')\n",
+    "        ax.set_title('Comparison of Target and Predicted Properties')\n",
+    "        ax.set_xticks(x)\n",
+    "        ax.set_xticklabels(labels, rotation=45, ha=\"right\")\n",
+    "        ax.legend()\n",
+    "\n",
+    "    except:\n",
+    "        print(\"Forward anaysis failed.\")\n",
+    "    return SMILES, is_novel\n",
+    "\n",
+    "def design_molecule_loop(model, tokenizer, target=None, temperature_gen=0.3,temperature_pred=0.01, SMILES_LIST=None,\n",
+    "                           top_k=50, top_p=0.95, repetition_penalty=1., num_beams=1,update_primer_with_better_draft=False,\n",
+    "                         threshold=0.01, N_max=100, dir_path='./',lower_bound = 0.0,remove_duplicates=True,\n",
+    "                upper_bound = 0.1,sample_count=0, messages=[], N_attempts_for_forward=1, set_opt=None):\n",
+    "\n",
+    "    mse_smallest_current=9999\n",
+    "    if not os.path.exists(dir_path):\n",
+    "        os.makedirs(dir_path)\n",
+    "    if target is None or not target.any():\n",
+    "        target = np.random.rand(12)\n",
+    "\n",
+    "    if len (messages) >0:\n",
+    "        print (\"Using primed generation:\\n\", messages)\n",
+    "    \n",
+    "    records = []  # To store SMILES, properties, and MSE\n",
+    "    for iteration in range(N_max):\n",
+    "        try:\n",
+    "            print (f\">>> Iteration={iteration}\")\n",
+    "            original_messages=copy.deepcopy (messages)\n",
+    "\n",
+    "            SMILES = design_from_target(model, tokenizer, target, temperature_gen, num_beams,\n",
+    "                                        top_k, top_p, repetition_penalty, messages=original_messages)\n",
+    "            is_novel=is_SMILES_novel (SMILES, SMILES_LIST)\n",
+    "\n",
+    "            if is_novel and is_valid_smiles(SMILES):\n",
+    "                print (f\"{SMILES} is novel: {is_novel}\", \"is valid: \", {is_valid_smiles(SMILES)})\n",
+    "                if N_attempts_for_forward==1:\n",
+    "                    predicted = properties_from_SMILES(model, tokenizer, SMILES,temperature_pred, num_beams,\n",
+    "                                                top_k, top_p, repetition_penalty)\n",
+    "                else:\n",
+    "                    predicted,_=avg_properties_from_SMILES(model, tokenizer, SMILES, SMILES_dir=dir_path,\n",
+    "                                    temperature=temperature_pred, top_k=top_k,top_p=top_p,    repetition_penalty=repetition_penalty,\n",
+    "                                    labels=labels, N_prop=N_attempts_for_forward, plot_results=False)\n",
+    "\n",
+    "                if set_opt==None:\n",
+    "                    mse = mean_squared_error(target, predicted)\n",
+    "                else:\n",
+    "                    mse = mean_squared_error(target[set_opt], predicted[set_opt])\n",
+    "                if mse<mse_smallest_current:\n",
+    "                    mse_smallest_current=mse\n",
+    "                    if update_primer_with_better_draft:\n",
+    "                        messages=prime_messages (SMILES, predicted , N=1)\n",
+    "                        print (\"Smaller MSE found, updated messages primer! Messages: \", messages,\n",
+    "                          f\"\\n\\nCurrent MSE: {mse}\")\n",
+    "                \n",
+    "                records.append((SMILES, predicted, mse, is_novel))\n",
+    "            \n",
+    "                print (f\">>>Iteration={iteration}, MSE={mse} for SMILES={SMILES}, novel={is_novel}\")\n",
+    "                if mse < threshold:\n",
+    "                    print(f\"Threshold met at iteration {iteration+1}\")\n",
+    "                    break\n",
+    "            else:\n",
+    "                print (f\"{SMILES} is not novel or not valid, validity: {is_valid_smiles(SMILES)}.\")\n",
+    "        except Exception as e:\n",
+    "            print(f\"Error during iteration {iteration+1}: {e}\")\n",
+    "            continue\n",
+    "\n",
+    "    # Sorting records based on MSE (most accurate first)\n",
+    "    records.sort(key=lambda x: x[2])\n",
+    "\n",
+    "    # Visualizing the best performing molecule\n",
+    "    best_SMILES, best_predicted, best_mse, is_novel = records[0]\n",
+    "\n",
+    "    print (\"Best SILES: \", best_SMILES)\n",
+    "    try:\n",
+    "        print (f\"{best_SMILES} is novel: {is_novel}\")\n",
+    "        \n",
+    "        sns.set_style(\"whitegrid\")\n",
+    "    \n",
+    "        visualize_pred_vs_target (target, best_predicted, labels, dir_path=dir_path, best_SMILES=best_SMILES,sample_count=0)\n",
+    "         \n",
+    "        print(f\"Process completed. Results saved to {csv_path}.\") \n",
+    "        visualize_SMILES(best_SMILES, dir_path=dir_path, root=f'{target}_BEST')\n",
+    "\n",
+    "        print(f\"Compute molecular structure, UFF eq, Gasteiger, etc.\") \n",
+    "        \n",
+    "        compute_gasteiger (best_SMILES, SMILES_dir=dir_path, target= np.array(best_predicted))\n",
+    "\n",
+    "        mol = Chem.MolFromSmiles(best_SMILES)\n",
+    "        inchi_str = Chem.MolToInchi(mol)\n",
+    "        print(f\"InChI String of {best_SMILES}:\", inchi_str)\n",
+    "    \n",
+    "           \n",
+    "    except Exception as e:\n",
+    "        print(f\"Processing/visualization failed for {best_SMILES}: {e}\")\n",
+    "\n",
+    "    # Writing records to a CSV file\n",
+    "    df = pd.DataFrame(records, columns=['SMILES', 'Predicted Properties', 'MSE', 'is_novel'])\n",
+    "    csv_path = os.path.join(dir_path, 'SMILES_designs.csv')\n",
+    "    df.to_csv(csv_path, index=False)\n",
+    "\n",
+    "    # Plot MSE against the index (which now corresponds to the ranking)\n",
+    "    plt.figure(figsize=(10, 8))  # Adjust the size as needed\n",
+    "    plt.plot(df['SMILES'], df['MSE'], 'o', markersize=5)  # 'o' for circular markers\n",
+    "    \n",
+    "    # Adding labels for each point with the SMILES string\n",
+    "    for i, txt in enumerate(df['SMILES']):\n",
+    "        plt.annotate(txt, (i, df['MSE'].iloc[i]), fontsize=8, rotation=45, ha='right')\n",
+    "    \n",
+    "    visualize_over_SMILES (df,N_max=N_max,SMILES_dir=SMILES_dir,\n",
+    "            lower_bound = lower_bound,remove_duplicates=remove_duplicates,\n",
+    "                upper_bound = upper_bound, target=target)\n",
+    "    return df \n",
+    "\n",
+    "from rdkit import Chem\n",
+    "from rdkit.Chem import Draw\n",
+    "import os\n",
+    "\n",
+    "def visualize_smiles_and_save(smiles_list, per_row=4, dir_path='./', root=''):\n",
+    "    \"\"\"\n",
+    "    Visualizes a list of molecules from their SMILES strings with labels, checks for validity, \n",
+    "    and saves the visualization as an SVG file.\n",
+    "    \n",
+    "    Parameters:\n",
+    "    - smiles_list: List of SMILES strings to visualize.\n",
+    "    - per_row: Number of molecule images per row in the assembly.\n",
+    "    - dir_path: Directory path where the SVG file will be saved.\n",
+    "    \"\"\"\n",
+    "    if not os.path.exists(dir_path):\n",
+    "        os.makedirs(dir_path)\n",
+    "    valid_molecules = []\n",
+    "    valid_smiles = []  # To store valid SMILES strings for labeling\n",
+    "    for smile in smiles_list:\n",
+    "        mol = Chem.MolFromSmiles(smile)\n",
+    "        if mol:  # If the molecule is valid\n",
+    "            valid_molecules.append(mol)\n",
+    "            valid_smiles.append(smile)  # Add the valid SMILES string\n",
+    "    \n",
+    "    # Proceed only if there are valid molecules\n",
+    "    if not valid_molecules:\n",
+    "        print(\"No valid molecules found in the provided SMILES strings.\")\n",
+    "        return\n",
+    "    \n",
+    "    # Ensure the directory exists\n",
+    "    if not os.path.exists(dir_path):\n",
+    "        os.makedirs(dir_path)\n",
+    "    \n",
+    "    # Define the SVG file path\n",
+    "    svg_file_path = os.path.join(dir_path, f'molecules_with_labels_{root}.svg')\n",
+    "    \n",
+    "    # Use RDKit to draw the molecules grid with labels\n",
+    "    fig = Draw.MolsToGridImage(valid_molecules, molsPerRow=per_row, subImgSize=(200, 200), \n",
+    "                               legends=valid_smiles, useSVG=True)\n",
+    "    \n",
+    "    # Saving the SVG content to a file\n",
+    "    with open(svg_file_path, 'w') as svg_file:\n",
+    "        svg_file.write(fig.data)\n",
+    "    display (fig)\n",
+    "    \n",
+    "    print(f\"Visualization saved as SVG at: {svg_file_path}\")\n",
+    "\n",
+    "    return valid_smiles \n",
+    "\n",
+    "def plot_MSE_over_SMILES (df_design,N_max=24,\n",
+    "            lower_bound = 0.0,\n",
+    "                upper_bound = 0.08, SMILES_dir='./', target='',  ):\n",
+    "    \n",
+    "    if not os.path.exists(SMILES_dir):\n",
+    "        os.makedirs(SMILES_dir)    \n",
+    "    df_sorted = df_design[:N_max].sort_values('MSE',ascending=False).reset_index(drop=True)\n",
+    "\n",
+    "    \n",
+    "    df_plot=df_sorted[(df_sorted['MSE'] > lower_bound) & (df_sorted['MSE'] < upper_bound)]\n",
+    "    \n",
+    "    # Plot MSE against the index (which now corresponds to the ranking)\n",
+    "    fig, ax = plt.subplots(figsize=(8, 7))\n",
+    "    plt.plot(df_plot['SMILES'], df_plot['MSE'], 'o-', markersize=5, )  # 'o' for circular markers\n",
+    "    \n",
+    "    # Improving the plot aesthetics\n",
+    "    plt.xticks(rotation=90)  # Rotate the x-axis labels for better readability\n",
+    "    plt.xlabel('Molecule SMILES')\n",
+    "    plt.ylabel('MSE')\n",
+    "    #plt.title('Ordered from Best to Worst')\n",
+    "    plt.tight_layout()  # Adjust the layout to make room for the rotated x-axis labels\n",
+    "    plt.savefig(SMILES_dir+f'SMILES_over_MSE_{target}.svg', format='svg')\n",
+    "    plt.show()\n",
+    "    \n",
+    "def visualize_over_SMILES (df_design,N_max=24,per_row=20,SMILES_dir='./',\n",
+    "            lower_bound = 0.0,\n",
+    "                upper_bound = 0.08, target='', remove_duplicates=True):\n",
+    "\n",
+    "    if remove_duplicates:\n",
+    "        # Example: Keep the entry with the best MSE among the novel molecules for each SMILES\n",
+    "        df_design = df_design.sort_values(['MSE', 'is_novel', 'SMILES', ], ascending=[True, False, True]) \\\n",
+    "             .drop_duplicates(subset='SMILES', keep='first')\n",
+    "\n",
+    "        df_design.reset_index(drop=True, inplace=True)\n",
+    "        df_design.to_csv(f'{SMILES_dir}/sorted_noduplicates_{N_max}.csv', index=False)\n",
+    "         \n",
+    "    valid_smiles=visualize_smiles_and_save(list(df_design['SMILES'][:N_max]), per_row=per_row, dir_path=SMILES_dir, root=f'{target}')\n",
+    "    \n",
+    "    smiles_df = pd.DataFrame(valid_smiles, columns=[\"SMILES\"])\n",
+    "\n",
+    "    # Save the DataFrame to a CSV file\n",
+    "    file_path = \"/smiles_data.csv\"\n",
+    "    smiles_df.to_csv(f'{SMILES_dir}/valid_SMILES_{N_max}.csv', index=False )\n",
+    "    \n",
+    "    fig, ax = plt.subplots(figsize=(8, 5))\n",
+    "    \n",
+    "    df_plot=df_design[(df_design['MSE'] > lower_bound) & (df_design['MSE'] < upper_bound)]\n",
+    "    df_plot.plot(kind='kde', color='darkblue', label='KDE', ax=ax)\n",
+    "    \n",
+    "    # Plot histogram with density=True for probability density representation\n",
+    "    plt.hist(df_design['MSE'], density=True, alpha=0.5, color='skyblue', label='Histogram',bins=50, \n",
+    "             range=[lower_bound,upper_bound]\n",
+    "            )\n",
+    "    plt.xlim(lower_bound, upper_bound)\n",
+    "    plt.title('Density and Histogram Plot of MSE')\n",
+    "    plt.xlabel('MSE')\n",
+    "    plt.ylabel('Density')\n",
+    "    \n",
+    "    # Adding a legend to distinguish between the KDE and Histogram\n",
+    "    plt.legend()\n",
+    "    \n",
+    "    plt.savefig(SMILES_dir+f'mse_histogram_{target}.svg', format='svg')\n",
+    "    plt.show()\n",
+    "\n",
+    "    plot_MSE_over_SMILES (df_design,N_max=N_max,\n",
+    "            lower_bound = lower_bound,\n",
+    "                upper_bound = upper_bound, target=target,SMILES_dir=SMILES_dir)\n",
+    "    \n",
+    "    return df_design\n",
+    "\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
+    "from pandas.plotting import parallel_coordinates\n",
+    "\n",
+    "def plot_change_in_design(original, labels, target, SMILES_dir='./'):\n",
+    "    if not os.path.exists(SMILES_dir):\n",
+    "        os.makedirs(SMILES_dir)\n",
+    "    \n",
+    "    # Create a DataFrame to hold the original and target vectors with labels\n",
+    "    df = pd.DataFrame([original, target], columns=labels)\n",
+    "    df['Version'] = ['Original', 'Target']  # Add a 'Version' column for coloring\n",
+    "    \n",
+    "    # Plotting\n",
+    "    plt.figure(figsize=(7, 4))\n",
+    "    parallel_coordinates(df, 'Version', color=['blue', 'red'])\n",
+    "    plt.title('Original vs Target Values across Properties')\n",
+    "    plt.xticks(rotation=45)\n",
+    "    plt.tight_layout()\n",
+    "    \n",
+    "    # Annotating changes with thicker arrows pointing towards the target\n",
+    "    for i, label in enumerate(labels):\n",
+    "        if original[i] < target[i]:  # If the target value is greater, arrow points upwards\n",
+    "            plt.annotate('', xy=(i, target[i]), xytext=(i, original[i]),\n",
+    "                         arrowprops=dict(arrowstyle=\"->\", color='black', lw=2))\n",
+    "        else:  # If the target value is lesser, arrow points downwards\n",
+    "            plt.annotate('', xy=(i, target[i]), xytext=(i, original[i]),\n",
+    "                         arrowprops=dict(arrowstyle=\"->\", color='black', lw=2))\n",
+    "    \n",
+    "    # Save the plot as an SVG file in the specified directory\n",
+    "    plt.savefig(SMILES_dir + \"parallel_coordinates_changes_direction.svg\", format=\"svg\")\n",
+    "    \n",
+    "    plt.show()\n",
+    "    \n",
+    "def visualize_pred_vs_target (target, best_predicted, labels, dir_path='./', best_SMILES='',sample_count=0):                           \n",
+    "    if not os.path.exists(dir_path):\n",
+    "        os.makedirs(dir_path)\n",
+    "    sns.set_style(\"whitegrid\")\n",
+    "    plt.gcf().set_facecolor('white')\n",
+    "    \n",
+    "    x = np.arange(len(labels))  # Label locations\n",
+    "    width = 0.35  # Width of the bars\n",
+    "    \n",
+    "    fig, ax = plt.subplots(figsize=(9, 5))\n",
+    "    rects1 = ax.bar(x - width/2, target, width, label='Target')\n",
+    "    rects2 = ax.bar(x + width/2, best_predicted, width, label='Predicted properties')\n",
+    "    \n",
+    "    # Add some text for labels, title and custom x-axis tick labels, etc.\n",
+    "    ax.set_ylabel('Values')\n",
+    "    ax.set_title(f'Comparison of Target and Predicted Properties, {best_SMILES}')\n",
+    "    ax.set_xticks(x)\n",
+    "    ax.set_xticklabels(labels, rotation=45, ha=\"right\")\n",
+    "    ax.legend()\n",
+    "    fig.tight_layout()\n",
+    "    plt.savefig(f\"{dir_path}/QM9_best_design_{target}_barplot_{sample_count}.svg\")\n",
+    "    plt.show()\n",
+    "    #plt.show()\n",
+    "\n",
+    "from rdkit import Chem\n",
+    "from rdkit.Chem import AllChem, Draw\n",
+    "from rdkit.Chem import AllChem, rdDepictor\n",
+    "from rdkit.Chem.Draw import rdMolDraw2D\n",
+    " \n",
+    "def prime_messages (SMILES_chitin_monomer, target, N=1):\n",
+    "    messages=[]\n",
+    "    for i in range (N):\n",
+    "        \n",
+    "        line=f'GenerateMolecularProperties<{return_str( target)}>'\n",
+    "        messages.append ({\"role\": \"user\", \"content\": line},  )\n",
+    "        line=f'[{SMILES_chitin_monomer}]'\n",
+    "        messages.append ({\"role\": \"assistant\", \"content\": line},  )\n",
+    "        \n",
+    "    return messages\n",
+    "\n",
+    "from rdkit import Chem\n",
+    "from rdkit.Chem import AllChem\n",
+    "\n",
+    "def smiles_to_3d(smiles, num_confs=100):\n",
+    "    mol = Chem.MolFromSmiles(smiles)\n",
+    "    if mol is None:\n",
+    "        print(\"Failed to create molecule from SMILES\")\n",
+    "        return None\n",
+    "\n",
+    "    mol = Chem.AddHs(mol)\n",
+    "    params = AllChem.ETKDGv3()\n",
+    "    params.randomSeed = 42\n",
+    "    if not AllChem.EmbedMultipleConfs(mol, numConfs=num_confs, params=params):\n",
+    "        print(\"Embedding conformations failed.\")\n",
+    "        return None\n",
+    "\n",
+    "    results = []\n",
+    "    for conf_id in range(num_confs):\n",
+    "        ff = AllChem.MMFFGetMoleculeForceField(mol, AllChem.MMFFGetMoleculeProperties(mol), confId=conf_id)\n",
+    "        if ff is None:\n",
+    "            print(f\"Failed to setup MMFF for conformer {conf_id}\")\n",
+    "            continue\n",
+    "        energy = ff.Minimize()\n",
+    "        results.append((conf_id, ff.CalcEnergy()))\n",
+    "\n",
+    "    if not results:\n",
+    "        print(\"No successful energy minimization.\")\n",
+    "        return None\n",
+    "    \n",
+    "\n",
+    "    best_conf = mol.GetConformer(min_energy_conf[0])\n",
+    "    best_mol = Chem.Mol(mol)\n",
+    "    best_mol.RemoveAllConformers()\n",
+    "    best_mol.AddConformer(best_conf, assignId=True)\n",
+    "\n",
+    "    coords = best_conf.GetPositions()\n",
+    "    atom_symbols = [atom.GetSymbol() for atom in best_mol.GetAtoms()]\n",
+    "    geometry = '\\n'.join(f'{atom} {coord[0]} {coord[1]} {coord[2]}' for atom, coord in zip(atom_symbols, coords))\n",
+    "\n",
+    "    display (best_mol)\n",
+    "    \n",
+    "    return geometry, best_mol"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "23f18039-4441-496c-89b0-9e467eaac83e",
+   "metadata": {},
+   "source": [
+    "### Property calculation as possible starting point for design iterations "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6519474d-4e03-4273-a79e-454d5845e6d6",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "SMILES_START='O1C2C3OC2C13'\n",
+    "properties,_=avg_properties_from_SMILES (model, tokenizer, SMILES_START, SMILES_dir=SMILES_dir,\n",
+    "                                temperature=0.3, top_k=256,top_p=0.9, num_beams=1,   repetition_penalty=1.,\n",
+    "                                labels=labels, N_prop=3, plot_results=True)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "198840ea-21f8-41eb-b62a-bc325261b731",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Retrieve the scaling parameters\n",
+    "data_min = scaler.data_min_\n",
+    "data_max = scaler.data_max_\n",
+    "scale = scaler.scale_\n",
+    "feature_min = scaler.min_\n",
+    "\n",
+    "print(\"Feature Scaling Parameters:\")\n",
+    "print(\"{:<20} {:<20} {:<20} {:<20}\".format(\"Feature Index\", \"Min Value\", \"Max Value\", \"Scale Factor\"))\n",
+    "for i in range(len(data_min)):\n",
+    "    print(\"{:<20} {:<20} {:<20} {:<20}\".format(i, data_min[i], data_max[i], scale[i]))\n",
+    "\n",
+    "print(\"\\nPer-feature Shifts (Min):\")\n",
+    "for i, min_val in enumerate(feature_min):\n",
+    "    print(\"Feature {}: {:.6f}\".format(i, min_val))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0dd1f217-74c0-40f3-8edc-4b610c12e0ea",
+   "metadata": {},
+   "source": [
+    "### Molecular design: Iterative solution "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc2747b6-90cc-4d42-bf93-bb39dc6d9198",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy \n",
+    "properties=y_test[4]\n",
+    "\n",
+    "#Create new set of properties based on existing molecule (from test set)\n",
+    "properties_new=copy.deepcopy (properties)\n",
+    "properties_new[0]=properties[0]+0.2\n",
+    "properties_new[1]=properties[1]+0.2\n",
+    "plot_change_in_design (properties, labels, properties_new,SMILES_dir)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c5f9b9c5-c746-48d1-841a-a2113d13279e",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "df_design=design_molecule_loop (model, tokenizer, np.array(properties_new), SMILES_LIST=SMILES_LIST, dir_path=SMILES_dir,\n",
+    "                          temperature_pred=0.1, temperature_gen=0.3, top_k=32,top_p=0.1,    repetition_penalty=1.,\n",
+    "                         threshold=0.001, N_max=64,  \n",
+    "                                N_attempts_for_forward=6,\n",
+    "                               )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8c5be323-aa47-49dd-bc44-be74936c62c8",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "visualize_over_SMILES (df_design,N_max=30,SMILES_dir=SMILES_dir,per_row=5,\n",
+    "            lower_bound = 0.0, remove_duplicates=True,\n",
+    "            upper_bound = 0.02, target=np.array(properties_new))\n",
+    "\n",
+    "target=np.array(properties_new)\n",
+    "best_SMILES, best_predicted, best_mse, is_novel = df_design_2.iloc[5]\n",
+    "\n",
+    "print (\"Best SILES: \", best_SMILES)\n",
+    "print (f\"{best_SMILES} is novel: {is_novel}\")\n",
+    "\n",
+    "sns.set_style(\"whitegrid\")\n",
+    "\n",
+    "visualize_pred_vs_target (target, best_predicted, labels, dir_path=SMILES_dir, best_SMILES=best_SMILES,sample_count=0)\n",
+    " \n",
+    "visualize_SMILES(best_SMILES, dir_path=SMILES_dir, root=f'{target}_BEST')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "25fd7dbe-95fd-4169-86e9-b05e86bbfb3a",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "target=np.array(properties_new)\n",
+    "best_SMILES, best_predicted, best_mse, is_novel = df_design_2.iloc[5]\n",
+    "\n",
+    "print (\"Best SILES: \", best_SMILES)\n",
+    "print (f\"{best_SMILES} is novel: {is_novel}\")\n",
+    "\n",
+    "sns.set_style(\"whitegrid\")\n",
+    "\n",
+    "visualize_pred_vs_target (target, best_predicted, labels, dir_path=SMILES_dir, best_SMILES=best_SMILES,sample_count=0)\n",
+    " \n",
+    "visualize_SMILES(best_SMILES, dir_path=SMILES_dir, root=f'{target}_BEST')"
+   ]
+  }
+ ],
+ "metadata": {
+  "environment": {
+   "kernel": "python3",
+   "name": ".m115",
+   "type": "gcloud",
+   "uri": "gcr.io/deeplearning-platform-release/:m115"
+  },
+  "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.11.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}