refactor

app.py

@@ -0,0 +1,288 @@
+import gradio as gr
+from inference import Inference
+import PIL
+from PIL import Image
+import pandas as pd
+import random
+from rdkit import Chem
+from rdkit.Chem import Draw
+from rdkit.Chem.Draw import IPythonConsole
+import shutil
+import os
+import time
+
+class DrugGENConfig:
+    # Inference configuration
+    submodel='DrugGEN'
+    inference_model="experiments/models/DrugGEN/"
+    sample_num=100
+    disable_correction=False  # corresponds to correct=True in old config
+
+    # Data configuration
+    inf_smiles='data/chembl_test.smi'  # corresponds to inf_raw_file in old config
+    train_smiles='data/chembl_train.smi'
+    train_drug_smiles='data/akt1_train.smi'
+    inf_batch_size=1
+    mol_data_dir='data'
+    features=False
+
+    # Model configuration
+    act='relu'
+    max_atom=45
+    dim=128
+    depth=1
+    heads=8
+    mlp_ratio=3
+    dropout=0.
+
+    # Seed configuration
+    set_seed=True
+    seed=10
+
+
+class DrugGENAKT1Config(DrugGENConfig):
+    submodel='DrugGEN'
+    inference_model="experiments/models/DrugGEN-AKT1/"
+    train_drug_smiles='data/akt1_train.smi'
+    max_atom=45
+
+
+class DrugGENCDK2Config(DrugGENConfig):
+    submodel='DrugGEN'
+    inference_model="experiments/models/DrugGEN-CDK2/"
+    train_drug_smiles='data/cdk2_train.smi'
+    max_atom=38
+
+
+class NoTargetConfig(DrugGENConfig):
+    submodel="NoTarget"
+    inference_model="experiments/models/NoTarget/"
+    train_drug_smiles='data/chembl_train.smi'  # No specific target, use general ChEMBL data
+
+
+model_configs = {
+    "DrugGEN-AKT1": DrugGENAKT1Config(),
+    "DrugGEN-CDK2": DrugGENCDK2Config(),
+    "DrugGEN-NoTarget": NoTargetConfig(),
+}
+
+
+
+def function(model_name: str, num_molecules: int, seed_num: int) -> tuple[PIL.Image, pd.DataFrame, str]:
+    '''
+    Returns:
+    image, score_df, file path
+    '''
+    if model_name == "DrugGEN-NoTarget":
+        model_name = "NoTarget"
+    
+    config = model_configs[model_name]
+    config.sample_num = num_molecules
+
+    if config.sample_num > 250:
+        raise gr.Error("You have requested to generate more than the allowed limit of 250 molecules. Please reduce your request to 250 or fewer.")
+
+    if seed_num is None or seed_num.strip() == "":
+        config.seed = random.randint(0, 10000)
+    else:
+        try:
+            config.seed = int(seed_num)
+        except ValueError:
+            raise gr.Error("The seed must be an integer value!")
+
+
+    inferer = Inference(config)
+    start_time = time.time()
+    scores = inferer.inference() # create scores_df out of this
+    et = time.time() - start_time
+
+    score_df = pd.DataFrame({
+        "Runtime (seconds)": [et],
+        "Validity": [scores["validity"].iloc[0]],
+        "Uniqueness": [scores["uniqueness"].iloc[0]],
+        "Novelty (Train)": [scores["novelty"].iloc[0]],
+        "Novelty (Test)": [scores["novelty_test"].iloc[0]],
+        "Drug Novelty": [scores["drug_novelty"].iloc[0]],
+        "Max Length": [scores["max_len"].iloc[0]],
+        "Mean Atom Type": [scores["mean_atom_type"].iloc[0]],
+        "SNN ChEMBL": [scores["snn_chembl"].iloc[0]],
+        "SNN Drug": [scores["snn_drug"].iloc[0]],
+        "Internal Diversity": [scores["IntDiv"].iloc[0]],
+        "QED": [scores["qed"].iloc[0]],
+        "SA Score": [scores["sa"].iloc[0]]
+    })
+
+    output_file_path = f'experiments/inference/{model_name}/inference_drugs.txt'
+
+    new_path = f'{model_name}_denovo_mols.smi'
+    os.rename(output_file_path, new_path)
+
+    with open(new_path) as f:
+        inference_drugs = f.read()
+
+    generated_molecule_list = inference_drugs.split("\n")[:-1]
+
+    rng = random.Random(config.seed)
+    if num_molecules > 12:
+        selected_molecules = rng.choices(generated_molecule_list, k=12)
+    else:
+        selected_molecules = generated_molecule_list
+    
+    selected_molecules = [Chem.MolFromSmiles(mol) for mol in selected_molecules if Chem.MolFromSmiles(mol) is not None]
+
+    drawOptions = Draw.rdMolDraw2D.MolDrawOptions()
+    drawOptions.prepareMolsBeforeDrawing = False
+    drawOptions.bondLineWidth = 0.5
+
+    molecule_image = Draw.MolsToGridImage(
+        selected_molecules,
+        molsPerRow=3,
+        subImgSize=(400, 400),
+        maxMols=len(selected_molecules),
+        # legends=None,
+        returnPNG=False,
+        drawOptions=drawOptions,
+        highlightAtomLists=None,
+        highlightBondLists=None,
+    )
+
+    return molecule_image, score_df, new_path
+
+
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("# DrugGEN: Target Centric De Novo Design of Drug Candidate Molecules with Graph Generative Deep Adversarial Networks")
+            with gr.Row():
+                gr.Markdown("[![arXiv](https://img.shields.io/badge/arXiv-2302.07868-b31b1b.svg)](https://arxiv.org/abs/2302.07868)")
+                gr.Markdown("[![github-repository](https://img.shields.io/badge/GitHub-black?logo=github)](https://github.com/HUBioDataLab/DrugGEN)")
+            
+            with gr.Accordion("About DrugGEN Models", open=False):
+                gr.Markdown("""
+## Model Variations
+
+### DrugGEN-AKT1
+This model is designed to generate molecules targeting the human AKT1 protein (UniProt ID: P31749), a serine/threonine-protein kinase that plays a key role in regulating cell survival, metabolism, and growth. AKT1 is a significant target in cancer therapy, particularly for breast, colorectal, and ovarian cancers.
+
+The model learns from:
+- General drug-like molecules from ChEMBL database
+- Known AKT1 inhibitors
+- Maximum atom count: 45
+
+### DrugGEN-CDK2
+This model targets the human CDK2 protein (UniProt ID: P24941), a cyclin-dependent kinase involved in cell cycle regulation. CDK2 inhibitors are being investigated for treating various cancers, particularly those with dysregulated cell cycle control.
+
+The model learns from:
+- General drug-like molecules from ChEMBL database
+- Known CDK2 inhibitors
+- Maximum atom count: 38
+
+### DrugGEN-NoTarget
+This is a general-purpose model that generates diverse drug-like molecules without targeting a specific protein. It's useful for:
+- Exploring chemical space
+- Generating diverse scaffolds
+- Creating molecules with drug-like properties
+
+## How It Works
+DrugGEN uses a graph-based generative adversarial network (GAN) architecture where:
+1. The generator creates molecular graphs
+2. The discriminator evaluates them against real molecules
+3. The model learns to generate increasingly realistic and target-specific molecules
+
+For more details, see our [paper on arXiv](https://arxiv.org/abs/2302.07868).
+                """)
+            
+            with gr.Accordion("Understanding the Metrics", open=False):
+                gr.Markdown("""
+## Evaluation Metrics
+
+### Basic Metrics
+- **Validity**: Percentage of generated molecules that are chemically valid
+- **Uniqueness**: Percentage of unique molecules among valid ones
+- **Runtime**: Time taken to generate the requested molecules
+
+### Novelty Metrics
+- **Novelty (Train)**: Percentage of molecules not found in the training set
+- **Novelty (Test)**: Percentage of molecules not found in the test set
+- **Drug Novelty**: Percentage of molecules not found in known drugs
+
+### Structural Metrics
+- **Max Length**: Maximum component length in the generated molecules
+- **Mean Atom Type**: Average distribution of atom types
+- **Internal Diversity**: Diversity within the generated set (higher is more diverse)
+
+### Drug-likeness Metrics
+- **QED (Quantitative Estimate of Drug-likeness)**: Score from 0-1 measuring how drug-like a molecule is (higher is better)
+- **SA Score (Synthetic Accessibility)**: Score from 1-10 indicating ease of synthesis (lower is easier)
+
+### Similarity Metrics
+- **SNN ChEMBL**: Similarity to ChEMBL molecules (higher means more similar to known drug-like compounds)
+- **SNN Drug**: Similarity to known drugs (higher means more similar to approved drugs)
+                """)
+            
+            model_name = gr.Radio(
+                choices=("DrugGEN-AKT1", "DrugGEN-CDK2", "DrugGEN-NoTarget"),
+                value="DrugGEN-AKT1",
+                label="Select Target Model",
+                info="Choose which protein target or general model to use for molecule generation"
+            )
+
+            num_molecules = gr.Slider(
+                minimum=10,
+                maximum=250,
+                value=100,
+                step=10,
+                label="Number of Molecules to Generate",
+                info="This space runs on a CPU, which may result in slower performance. Generating 200 molecules takes approximately 6 minutes. Therefore, We set a 250-molecule cap. On a GPU, the model can generate 10,000 molecules in the same amount of time. Please check our GitHub repo for running our models on GPU.""
+            )
+
+            seed_num = gr.Textbox(
+                label="Random Seed (Optional)",
+                value="",
+                info="Set a specific seed for reproducible results, or leave empty for random generation"
+            )
+
+            submit_button = gr.Button(
+                value="Generate Molecules",
+                variant="primary",
+                size="lg"
+            )
+
+        with gr.Column(scale=2):
+            with gr.Tabs():
+                with gr.TabItem("Generated Molecules"):
+                    image_output = gr.Image(
+                        label="Sample of Generated Molecules",
+                        elem_id="molecule_display"
+                    )
+                    file_download = gr.File(
+                        label="Download All Generated Molecules (SMILES format)",
+                    )
+                
+                with gr.TabItem("Performance Metrics"):
+                    scores_df = gr.Dataframe(
+                        label="Model Performance Metrics",
+                        headers=["Runtime (seconds)", "Validity", "Uniqueness", "Novelty (Train)", "Novelty (Test)", 
+                                "Drug Novelty", "Max Length", "Mean Atom Type", "SNN ChEMBL", "SNN Drug", 
+                                "Internal Diversity", "QED", "SA Score"]
+                    )
+            
+            with gr.Accordion("Generation Settings", open=False):
+                gr.Markdown("""
+                ## Technical Details
+                
+                - This demo runs on CPU which limits generation speed
+                - Generating 200 molecules takes approximately 6 minutes
+                - For faster generation or larger batches, run the model on GPU using our GitHub repository
+                - The model uses a graph-based representation of molecules
+                - Maximum atom count varies by model (AKT1: 45, CDK2: 38)
+                """)
+
+    gr.Markdown("### Created by the HU BioDataLab | [GitHub](https://github.com/HUBioDataLab/DrugGEN) | [Paper](https://arxiv.org/abs/2302.07868)")
+
+    submit_button.click(function, inputs=[model_name, num_molecules, seed_num], outputs=[image_output, scores_df, file_download], api_name="inference")
+#demo.queue(concurrency_count=1)
+demo.queue()
+demo.launch()
+

inference.py

@@ -0,0 +1,303 @@
+import os
+import sys
+import time
+import random
+import pickle
+import argparse
+import os.path as osp
+
+import torch
+import torch.utils.data
+from torch_geometric.loader import DataLoader
+
+import pandas as pd
+from tqdm import tqdm
+
+from rdkit import RDLogger, Chem
+from rdkit.Chem import QED, RDConfig
+
+sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
+import sascorer
+
+from src.util.utils import *
+from src.model.models import Generator
+from src.data.dataset import DruggenDataset
+from src.data.utils import get_encoders_decoders, load_molecules
+from src.model.loss import generator_loss
+from src.util.smiles_cor import smi_correct
+
+
+class Inference(object):
+    """Inference class for DrugGEN."""
+
+    def __init__(self, config):
+        if config.set_seed:
+            np.random.seed(config.seed)
+            random.seed(config.seed)
+            torch.manual_seed(config.seed)
+            torch.cuda.manual_seed_all(config.seed)
+
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+            os.environ["PYTHONHASHSEED"] = str(config.seed)
+
+            print(f'Using seed {config.seed}')
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
+
+        # Initialize configurations
+        self.submodel = config.submodel
+        self.inference_model = config.inference_model
+        self.sample_num = config.sample_num
+        self.disable_correction = config.disable_correction
+
+        # Data loader.
+        self.inf_smiles = config.inf_smiles  # SMILES containing text file for first dataset. 
+                                         # Write the full path to file.
+        
+        inf_smiles_basename = osp.basename(self.inf_smiles)
+        
+        # Get the base name without extension and add max_atom to it
+        self.max_atom = config.max_atom  # Model is based on one-shot generation.
+        inf_smiles_base = os.path.splitext(inf_smiles_basename)[0]
+        
+        # Change extension from .smi to .pt and add max_atom to the filename
+        self.inf_dataset_file = f"{inf_smiles_base}{self.max_atom}.pt"
+
+        self.inf_batch_size = config.inf_batch_size
+        self.train_smiles = config.train_smiles
+        self.train_drug_smiles = config.train_drug_smiles
+        self.mol_data_dir = config.mol_data_dir  # Directory where the dataset files are stored.
+        self.dataset_name = self.inf_dataset_file.split(".")[0]
+        self.features = config.features  # Small model uses atom types as node features. (Boolean, False uses atom types only.)
+                                         # Additional node features can be added. Please check new_dataloarder.py Line 102.
+
+        # Get atom and bond encoders/decoders
+        self.atom_encoder, self.atom_decoder, self.bond_encoder, self.bond_decoder = get_encoders_decoders(
+            self.train_smiles,
+            self.train_drug_smiles,
+            self.max_atom
+        )
+
+        self.inf_dataset = DruggenDataset(self.mol_data_dir,
+                                      self.inf_dataset_file,
+                                      self.inf_smiles,
+                                      self.max_atom,
+                                      self.features,
+                                      atom_encoder=self.atom_encoder,
+                                      atom_decoder=self.atom_decoder,
+                                      bond_encoder=self.bond_encoder,
+                                      bond_decoder=self.bond_decoder)
+
+        self.inf_loader = DataLoader(self.inf_dataset,
+                                 shuffle=True,
+                                 batch_size=self.inf_batch_size,
+                                 drop_last=True)  # PyG dataloader for the first GAN.
+
+        self.m_dim = len(self.atom_decoder) if not self.features else int(self.inf_loader.dataset[0].x.shape[1]) # Atom type dimension.
+        self.b_dim = len(self.bond_decoder) # Bond type dimension.
+        self.vertexes = int(self.inf_loader.dataset[0].x.shape[0]) # Number of nodes in the graph.
+
+        # Model configurations.
+        self.act = config.act
+        self.dim = config.dim
+        self.depth = config.depth
+        self.heads = config.heads
+        self.mlp_ratio = config.mlp_ratio
+        self.dropout = config.dropout
+
+        self.build_model()
+
+    def build_model(self):
+        """Create generators and discriminators."""
+        self.G = Generator(self.act,
+                           self.vertexes,
+                           self.b_dim,
+                           self.m_dim,
+                           self.dropout,
+                           dim=self.dim,
+                           depth=self.depth,
+                           heads=self.heads,
+                           mlp_ratio=self.mlp_ratio)
+        self.G.to(self.device)
+        self.print_network(self.G, 'G')
+
+    def print_network(self, model, name):
+        """Print out the network information."""
+        num_params = 0
+        for p in model.parameters():
+            num_params += p.numel() 
+        print(model)
+        print(name)
+        print("The number of parameters: {}".format(num_params))
+
+    def restore_model(self, submodel, model_directory):
+        """Restore the trained generator and discriminator."""
+        print('Loading the model...')
+        G_path = os.path.join(model_directory, '{}-G.ckpt'.format(submodel))
+        self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
+
+    def inference(self):
+        # Load the trained generator.
+        self.restore_model(self.submodel, self.inference_model)
+
+        # smiles data for metrics calculation.
+        chembl_smiles = [line for line in open(self.train_smiles, 'r').read().splitlines()]
+        chembl_test = [line for line in open(self.inf_smiles, 'r').read().splitlines()]
+        drug_smiles = [line for line in open(self.train_drug_smiles, 'r').read().splitlines()]
+        drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
+        drug_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in drug_mols if x is not None]
+
+
+        # Make directories if not exist.
+        if not os.path.exists("experiments/inference/{}".format(self.submodel)):
+            os.makedirs("experiments/inference/{}".format(self.submodel))
+
+        if not self.disable_correction:
+            correct = smi_correct(self.submodel, "experiments/inference/{}".format(self.submodel))
+
+        search_res = pd.DataFrame(columns=["submodel", "validity",
+                                           "uniqueness", "novelty",
+                                           "novelty_test", "drug_novelty",
+                                           "max_len", "mean_atom_type",
+                                           "snn_chembl", "snn_drug", "IntDiv", "qed", "sa"])
+
+        self.G.eval()
+
+        start_time = time.time()
+        metric_calc_dr = []
+        uniqueness_calc = []
+        real_smiles_snn = []
+        nodes_sample = torch.Tensor(size=[1, self.vertexes, 1]).to(self.device)
+        generated_smiles = []
+        val_counter = 0
+        none_counter = 0
+
+        # Inference mode
+        with torch.inference_mode():
+            pbar = tqdm(range(self.sample_num))
+            pbar.set_description('Inference mode for {} model started'.format(self.submodel))
+            for i, data in enumerate(self.inf_loader):
+                val_counter += 1
+                # Preprocess dataset 
+                _, a_tensor, x_tensor = load_molecules(
+                    data=data, 
+                    batch_size=self.inf_batch_size,
+                    device=self.device,
+                    b_dim=self.b_dim,
+                    m_dim=self.m_dim,
+                )
+
+                _, _, node_sample, edge_sample = self.G(a_tensor, x_tensor)
+
+                g_edges_hat_sample = torch.max(edge_sample, -1)[1]
+                g_nodes_hat_sample = torch.max(node_sample, -1)[1]
+
+                fake_mol_g = [self.inf_dataset.matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=False, file_name=self.dataset_name) 
+                        for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
+
+                a_tensor_sample = torch.max(a_tensor, -1)[1]
+                x_tensor_sample = torch.max(x_tensor, -1)[1]
+                real_mols = [self.inf_dataset.matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(), strict=True, file_name=self.dataset_name) 
+                        for e_, n_ in zip(a_tensor_sample, x_tensor_sample)]
+
+                inference_drugs = [None if line is None else Chem.MolToSmiles(line) for line in fake_mol_g]
+                inference_drugs = [None if x is None else max(x.split('.'), key=len) for x in inference_drugs]
+
+                for molecules in inference_drugs:
+                    if molecules is None:
+                        none_counter += 1
+
+                for molecules in inference_drugs:
+                    if molecules is not None:
+                        molecules = molecules.replace("*", "C")
+                        generated_smiles.append(molecules)
+                        uniqueness_calc.append(molecules)
+                        nodes_sample = torch.cat((nodes_sample, g_nodes_hat_sample.view(1, self.vertexes, 1)), 0)
+                        pbar.update(1)
+                    metric_calc_dr.append(molecules)
+
+                real_smiles_snn.append(real_mols[0])
+                generation_number = len([x for x in metric_calc_dr if x is not None])
+                if generation_number == self.sample_num or none_counter == self.sample_num:
+                    break
+
+        if not self.disable_correction:
+            correct = smi_correct(self.submodel, "experiments/inference/{}".format(self.submodel))
+            gen_smi = correct.correct_smiles_list(generated_smiles)
+        else:
+            gen_smi = generated_smiles
+
+        et = time.time() - start_time
+        
+        gen_vecs = [AllChem.GetMorganFingerprintAsBitVect(Chem.MolFromSmiles(x), 2, nBits=1024) for x in uniqueness_calc if Chem.MolFromSmiles(x) is not None]
+        real_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in real_smiles_snn if x is not None]
+
+        if not self.disable_correction:
+            val = round(len(gen_smi)/self.sample_num, 3)
+        else: 
+            val = round(fraction_valid(gen_smi), 3)
+
+        uniq = round(fraction_unique(gen_smi), 3)
+        nov = round(novelty(gen_smi, chembl_smiles), 3)
+        nov_test = round(novelty(gen_smi, chembl_test), 3)
+        drug_nov = round(novelty(gen_smi, drug_smiles), 3)
+        max_len = round(Metrics.max_component(gen_smi, self.vertexes), 3)
+        mean_atom = round(Metrics.mean_atom_type(nodes_sample), 3)
+        snn_chembl = round(average_agg_tanimoto(np.array(real_vecs), np.array(gen_vecs)), 3)
+        snn_drug = round(average_agg_tanimoto(np.array(drug_vecs), np.array(gen_vecs)), 3)
+        int_div = round((internal_diversity(np.array(gen_vecs)))[0], 3)
+        qed = round(np.mean([QED.qed(Chem.MolFromSmiles(x)) for x in gen_smi if Chem.MolFromSmiles(x) is not None]), 3)
+        sa = round(np.mean([sascorer.calculateScore(Chem.MolFromSmiles(x)) for x in gen_smi if Chem.MolFromSmiles(x) is not None]), 3)
+
+        model_res = pd.DataFrame({"submodel": [self.submodel], "validity": [val],
+                        "uniqueness": [uniq], "novelty": [nov],
+                        "novelty_test": [nov_test], "drug_novelty": [drug_nov],
+                        "max_len": [max_len], "mean_atom_type": [mean_atom],
+                        "snn_chembl": [snn_chembl], "snn_drug": [snn_drug], 
+                        "IntDiv": [int_div], "qed": [qed], "sa": [sa]})
+
+        # Write generated SMILES to a temporary file for app.py to use
+        temp_file = f'{self.submodel}_denovo_mols.smi'
+        with open(temp_file, 'w') as f:
+            f.write("SMILES\n")
+            for smiles in gen_smi:
+                f.write(f"{smiles}\n")
+
+        return model_res
+
+
+if __name__=="__main__":
+    parser = argparse.ArgumentParser()
+
+    # Inference configuration.
+    parser.add_argument('--submodel', type=str, default="DrugGEN", help="Chose model subtype: DrugGEN, NoTarget", choices=['DrugGEN', 'NoTarget'])
+    parser.add_argument('--inference_model', type=str, help="Path to the model for inference")
+    parser.add_argument('--sample_num', type=int, default=100, help='inference samples')
+    parser.add_argument('--disable_correction', action='store_true', help='Disable SMILES correction')
+   
+    # Data configuration.
+    parser.add_argument('--inf_smiles', type=str, required=True)
+    parser.add_argument('--train_smiles', type=str, required=True)
+    parser.add_argument('--train_drug_smiles', type=str, required=True)
+    parser.add_argument('--inf_batch_size', type=int, default=1, help='Batch size for inference')
+    parser.add_argument('--mol_data_dir', type=str, default='data')
+    parser.add_argument('--features', action='store_true', help='features dimension for nodes')
+
+    # Model configuration.
+    parser.add_argument('--act', type=str, default="relu", help="Activation function for the model.", choices=['relu', 'tanh', 'leaky', 'sigmoid'])
+    parser.add_argument('--max_atom', type=int, default=45, help='Max atom number for molecules must be specified.')
+    parser.add_argument('--dim', type=int, default=128, help='Dimension of the Transformer Encoder model for the GAN.')
+    parser.add_argument('--depth', type=int, default=1, help='Depth of the Transformer model from the GAN.')
+    parser.add_argument('--heads', type=int, default=8, help='Number of heads for the MultiHeadAttention module from the GAN.')
+    parser.add_argument('--mlp_ratio', type=int, default=3, help='MLP ratio for the Transformer.')
+    parser.add_argument('--dropout', type=float, default=0., help='dropout rate')
+
+    # Seed configuration.
+    parser.add_argument('--set_seed', action='store_true', help='set seed for reproducibility')
+    parser.add_argument('--seed', type=int, default=1, help='seed for reproducibility')
+
+    config = parser.parse_args()
+    inference = Inference(config)
+    inference.inference()

src/data/dataset.py

@@ -0,0 +1,317 @@
+import os
+import os.path as osp
+import re
+import pickle
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+import torch
+from torch_geometric.data import Data, InMemoryDataset
+
+from rdkit import Chem, RDLogger
+
+from src.data.utils import label2onehot
+
+RDLogger.DisableLog('rdApp.*') 
+
+
+class DruggenDataset(InMemoryDataset):
+    def __init__(self, root, dataset_file, raw_files, max_atom, features, 
+                 atom_encoder, atom_decoder, bond_encoder, bond_decoder,
+                 transform=None, pre_transform=None, pre_filter=None):
+        """
+        Initialize the DruggenDataset with pre-loaded encoder/decoder dictionaries.
+
+        Parameters:
+            root (str): Root directory.
+            dataset_file (str): Name of the processed dataset file.
+            raw_files (str): Path to the raw SMILES file.
+            max_atom (int): Maximum number of atoms allowed in a molecule.
+            features (bool): Whether to include additional node features.
+            atom_encoder (dict): Pre-loaded atom encoder dictionary.
+            atom_decoder (dict): Pre-loaded atom decoder dictionary.
+            bond_encoder (dict): Pre-loaded bond encoder dictionary.
+            bond_decoder (dict): Pre-loaded bond decoder dictionary.
+            transform, pre_transform, pre_filter: See PyG InMemoryDataset.
+        """
+        self.dataset_name = dataset_file.split(".")[0]
+        self.dataset_file = dataset_file
+        self.raw_files = raw_files
+        self.max_atom = max_atom
+        self.features = features
+
+        # Use the provided encoder/decoder mappings.
+        self.atom_encoder_m = atom_encoder
+        self.atom_decoder_m = atom_decoder
+        self.bond_encoder_m = bond_encoder
+        self.bond_decoder_m = bond_decoder
+
+        self.atom_num_types = len(atom_encoder)
+        self.bond_num_types = len(bond_encoder)
+
+        super().__init__(root, transform, pre_transform, pre_filter)
+        path = osp.join(self.processed_dir, dataset_file)
+        self.data, self.slices = torch.load(path)
+        self.root = root
+
+    @property
+    def processed_dir(self):
+        """
+        Returns the directory where processed dataset files are stored.
+        """
+        return self.root
+    
+    @property
+    def raw_file_names(self):
+        """
+        Returns the raw SMILES file name.
+        """
+        return self.raw_files
+
+    @property
+    def processed_file_names(self):
+        """
+        Returns the name of the processed dataset file.
+        """
+        return self.dataset_file
+
+    def _filter_smiles(self, smiles_list):
+        """
+        Filters the input list of SMILES strings to keep only valid molecules that:
+         - Can be successfully parsed,
+         - Have a number of atoms less than or equal to the maximum allowed (max_atom),
+         - Contain only atoms present in the atom_encoder,
+         - Contain only bonds present in the bond_encoder.
+
+        Parameters:
+            smiles_list (list): List of SMILES strings.
+
+        Returns:
+            max_length (int): Maximum number of atoms found in the filtered molecules.
+            filtered_smiles (list): List of valid SMILES strings.
+        """
+        max_length = 0
+        filtered_smiles = []
+        for smiles in tqdm(smiles_list, desc="Filtering SMILES"):
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                continue
+
+            # Check molecule size
+            molecule_size = mol.GetNumAtoms()
+            if molecule_size > self.max_atom:
+                continue
+
+            # Filter out molecules with atoms not in the atom_encoder
+            if not all(atom.GetAtomicNum() in self.atom_encoder_m for atom in mol.GetAtoms()):
+                continue
+
+            # Filter out molecules with bonds not in the bond_encoder
+            if not all(bond.GetBondType() in self.bond_encoder_m for bond in mol.GetBonds()):
+                continue
+
+            filtered_smiles.append(smiles)
+            max_length = max(max_length, molecule_size)
+        return max_length, filtered_smiles
+
+    def _genA(self, mol, connected=True, max_length=None):
+        """
+        Generates the adjacency matrix for a molecule based on its bond structure.
+
+        Parameters:
+            mol (rdkit.Chem.Mol): The molecule.
+            connected (bool): If True, ensures all atoms are connected.
+            max_length (int, optional): The size of the matrix; if None, uses number of atoms in mol.
+
+        Returns:
+            np.array: Adjacency matrix with bond types as entries, or None if disconnected.
+        """
+        max_length = max_length if max_length is not None else mol.GetNumAtoms()
+        A = np.zeros((max_length, max_length))
+        begin = [b.GetBeginAtomIdx() for b in mol.GetBonds()]
+        end = [b.GetEndAtomIdx() for b in mol.GetBonds()]
+        bond_type = [self.bond_encoder_m[b.GetBondType()] for b in mol.GetBonds()]
+        A[begin, end] = bond_type
+        A[end, begin] = bond_type
+        degree = np.sum(A[:mol.GetNumAtoms(), :mol.GetNumAtoms()], axis=-1)
+        return A if connected and (degree > 0).all() else None
+
+    def _genX(self, mol, max_length=None):
+        """
+        Generates the feature vector for each atom in a molecule by encoding their atomic numbers.
+
+        Parameters:
+            mol (rdkit.Chem.Mol): The molecule.
+            max_length (int, optional): Length of the feature vector; if None, uses number of atoms in mol.
+
+        Returns:
+            np.array: Array of atom feature indices, padded with zeros if necessary, or None on error.
+        """
+        max_length = max_length if max_length is not None else mol.GetNumAtoms()
+        try:
+            return np.array([self.atom_encoder_m[atom.GetAtomicNum()] for atom in mol.GetAtoms()] +
+                            [0] * (max_length - mol.GetNumAtoms()))
+        except KeyError as e:
+            print(f"Skipping molecule with unsupported atom: {e}")
+            print(f"Skipped SMILES: {Chem.MolToSmiles(mol)}")
+            return None
+
+    def _genF(self, mol, max_length=None):
+        """
+        Generates additional node features for a molecule using various atomic properties.
+
+        Parameters:
+            mol (rdkit.Chem.Mol): The molecule.
+            max_length (int, optional): Number of rows in the features matrix; if None, uses number of atoms.
+
+        Returns:
+            np.array: Array of additional features for each atom, padded with zeros if necessary.
+        """
+        max_length = max_length if max_length is not None else mol.GetNumAtoms()
+        features = np.array([[*[a.GetDegree() == i for i in range(5)],
+                               *[a.GetExplicitValence() == i for i in range(9)],
+                               *[int(a.GetHybridization()) == i for i in range(1, 7)],
+                               *[a.GetImplicitValence() == i for i in range(9)],
+                               a.GetIsAromatic(),
+                               a.GetNoImplicit(),
+                               *[a.GetNumExplicitHs() == i for i in range(5)],
+                               *[a.GetNumImplicitHs() == i for i in range(5)],
+                               *[a.GetNumRadicalElectrons() == i for i in range(5)],
+                               a.IsInRing(),
+                               *[a.IsInRingSize(i) for i in range(2, 9)]]
+                              for a in mol.GetAtoms()], dtype=np.int32)
+        return np.vstack((features, np.zeros((max_length - features.shape[0], features.shape[1]))))
+
+    def decoder_load(self, dictionary_name, file):
+        """
+        Returns the pre-loaded decoder dictionary based on the dictionary name.
+
+        Parameters:
+            dictionary_name (str): Name of the dictionary ("atom" or "bond").
+            file: Placeholder parameter for compatibility.
+
+        Returns:
+            dict: The corresponding decoder dictionary.
+        """
+        if dictionary_name == "atom":
+            return self.atom_decoder_m
+        elif dictionary_name == "bond":
+            return self.bond_decoder_m
+        else:
+            raise ValueError("Unknown dictionary name.")
+
+    def matrices2mol(self, node_labels, edge_labels, strict=True, file_name=None):
+        """
+        Converts graph representations (node labels and edge labels) back to an RDKit molecule.
+
+        Parameters:
+            node_labels (iterable): Encoded atom labels.
+            edge_labels (np.array): Adjacency matrix with encoded bond types.
+            strict (bool): If True, sanitizes the molecule and returns None on failure.
+            file_name: Placeholder parameter for compatibility.
+
+        Returns:
+            rdkit.Chem.Mol: The resulting molecule, or None if sanitization fails.
+        """
+        mol = Chem.RWMol()
+        for node_label in node_labels:
+            mol.AddAtom(Chem.Atom(self.atom_decoder_m[node_label]))
+        for start, end in zip(*np.nonzero(edge_labels)):
+            if start > end:
+                mol.AddBond(int(start), int(end), self.bond_decoder_m[edge_labels[start, end]])
+        if strict:
+            try:
+                Chem.SanitizeMol(mol)
+            except Exception:
+                mol = None
+        return mol
+
+    def check_valency(self, mol):
+        """
+        Checks that no atom in the molecule has exceeded its allowed valency.
+
+        Parameters:
+            mol (rdkit.Chem.Mol): The molecule.
+
+        Returns:
+            tuple: (True, None) if valid; (False, atomid_valence) if there is a valency issue.
+        """
+        try:
+            Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_PROPERTIES)
+            return True, None
+        except ValueError as e:
+            e = str(e)
+            p = e.find('#')
+            e_sub = e[p:]
+            atomid_valence = list(map(int, re.findall(r'\d+', e_sub)))
+            return False, atomid_valence
+
+    def correct_mol(self, mol):
+        """
+        Corrects a molecule by removing bonds until all atoms satisfy their valency limits.
+
+        Parameters:
+            mol (rdkit.Chem.Mol): The molecule.
+
+        Returns:
+            rdkit.Chem.Mol: The corrected molecule.
+        """
+        while True:
+            flag, atomid_valence = self.check_valency(mol)
+            if flag:
+                break
+            else:
+                # Expecting two numbers: atom index and its valence.
+                assert len(atomid_valence) == 2
+                idx = atomid_valence[0]
+                queue = []
+                for b in mol.GetAtomWithIdx(idx).GetBonds():
+                    queue.append((b.GetIdx(), int(b.GetBondType()), b.GetBeginAtomIdx(), b.GetEndAtomIdx()))
+                queue.sort(key=lambda tup: tup[1], reverse=True)
+                if queue:
+                    start = queue[0][2]
+                    end = queue[0][3]
+                    mol.RemoveBond(start, end)
+        return mol
+
+
+    def process(self, size=None):
+        """
+        Processes the raw SMILES file by filtering and converting each valid SMILES into a PyTorch Geometric Data object.
+        The resulting dataset is saved to disk.
+
+        Parameters:
+            size (optional): Placeholder parameter for compatibility.
+
+        Side Effects:
+            Saves the processed dataset as a file in the processed directory.
+        """
+        # Read raw SMILES from file (assuming CSV with no header)
+        smiles_list = pd.read_csv(self.raw_files, header=None)[0].tolist()
+        max_length, filtered_smiles = self._filter_smiles(smiles_list)
+        data_list = []
+        self.m_dim = len(self.atom_decoder_m)
+        for smiles in tqdm(filtered_smiles, desc='Processing dataset', total=len(filtered_smiles)):
+            mol = Chem.MolFromSmiles(smiles)
+            A = self._genA(mol, connected=True, max_length=max_length)
+            if A is not None:
+                x_array = self._genX(mol, max_length=max_length)
+                if x_array is None:
+                    continue
+                x = torch.from_numpy(x_array).to(torch.long).view(1, -1)
+                x = label2onehot(x, self.m_dim).squeeze()
+                if self.features:
+                    f = torch.from_numpy(self._genF(mol, max_length=max_length)).to(torch.long).view(x.shape[0], -1)
+                    x = torch.concat((x, f), dim=-1)
+                adjacency = torch.from_numpy(A)
+                edge_index = adjacency.nonzero(as_tuple=False).t().contiguous()
+                edge_attr = adjacency[edge_index[0], edge_index[1]].to(torch.long)
+                data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, smiles=smiles)
+                if self.pre_filter is not None and not self.pre_filter(data):
+                    continue
+                if self.pre_transform is not None:
+                    data = self.pre_transform(data)
+                data_list.append(data)
+        torch.save(self.collate(data_list), osp.join(self.processed_dir, self.dataset_file))

src/data/utils.py

@@ -0,0 +1,143 @@
+import os
+import pickle
+
+import pandas as pd
+from tqdm import tqdm
+
+import torch
+from torch_geometric.data import Data, InMemoryDataset
+import torch_geometric.utils as geoutils
+
+from rdkit import Chem, RDLogger
+
+
+
+def label2onehot(labels, dim, device=None):
+    """Convert label indices to one-hot vectors."""
+    out = torch.zeros(list(labels.size())+[dim])
+    if device:
+        out = out.to(device)
+
+    out.scatter_(len(out.size())-1,labels.unsqueeze(-1),1.)
+
+    return out.float()
+
+
+def get_encoders_decoders(raw_file1, raw_file2, max_atom):
+    """
+    Given two raw SMILES files, either load the atom and bond encoders/decoders
+    if they exist (naming them based on the file names) or create and save them.
+
+    Parameters:
+        raw_file1 (str): Path to the first SMILES file.
+        raw_file2 (str): Path to the second SMILES file.
+        max_atom (int): Maximum allowed number of atoms in a molecule.
+
+    Returns:
+        atom_encoder (dict): Mapping from atomic numbers to indices.
+        atom_decoder (dict): Mapping from indices to atomic numbers.
+        bond_encoder (dict): Mapping from bond types to indices.
+        bond_decoder (dict): Mapping from indices to bond types.
+    """
+    # Determine unique suffix based on the two file names (alphabetically sorted for consistency)
+    name1 = os.path.splitext(os.path.basename(raw_file1))[0]
+    name2 = os.path.splitext(os.path.basename(raw_file2))[0]
+    sorted_names = sorted([name1, name2])
+    suffix = f"{sorted_names[0]}_{sorted_names[1]}"
+
+    # Define encoder/decoder directories and file paths
+    enc_dir = os.path.join("data", "encoders")
+    dec_dir = os.path.join("data", "decoders")
+    atom_encoder_path = os.path.join(enc_dir, f"atom_{suffix}.pkl")
+    atom_decoder_path = os.path.join(dec_dir, f"atom_{suffix}.pkl")
+    bond_encoder_path = os.path.join(enc_dir, f"bond_{suffix}.pkl")
+    bond_decoder_path = os.path.join(dec_dir, f"bond_{suffix}.pkl")
+
+    # If all files exist, load and return them
+    if (os.path.exists(atom_encoder_path) and os.path.exists(atom_decoder_path) and 
+        os.path.exists(bond_encoder_path) and os.path.exists(bond_decoder_path)):
+        with open(atom_encoder_path, "rb") as f:
+            atom_encoder = pickle.load(f)
+        with open(atom_decoder_path, "rb") as f:
+            atom_decoder = pickle.load(f)
+        with open(bond_encoder_path, "rb") as f:
+            bond_encoder = pickle.load(f)
+        with open(bond_decoder_path, "rb") as f:
+            bond_decoder = pickle.load(f)
+        print("Loaded existing encoders/decoders!")
+        return atom_encoder, atom_decoder, bond_encoder, bond_decoder
+
+    # Otherwise, create the encoders/decoders
+    print("Creating new encoders/decoders...")
+    # Read SMILES from both files (assuming one SMILES per row, no header)
+    smiles1 = pd.read_csv(raw_file1, header=None)[0].tolist()
+    smiles2 = pd.read_csv(raw_file2, header=None)[0].tolist()
+    smiles_combined = smiles1 + smiles2
+
+    atom_labels = set()
+    bond_labels = set()
+    max_length = 0
+    filtered_smiles = []
+    
+    # Process each SMILES: keep only valid molecules with <= max_atom atoms
+    for smiles in tqdm(smiles_combined, desc="Processing SMILES"):
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            continue
+        molecule_size = mol.GetNumAtoms()
+        if molecule_size > max_atom:
+            continue
+        filtered_smiles.append(smiles)
+        # Collect atomic numbers
+        atom_labels.update([atom.GetAtomicNum() for atom in mol.GetAtoms()])
+        max_length = max(max_length, molecule_size)
+        # Collect bond types
+        bond_labels.update([bond.GetBondType() for bond in mol.GetBonds()])
+    
+    # Add a PAD symbol (here using 0 for atoms)
+    atom_labels.add(0)
+    atom_labels = sorted(atom_labels)
+    
+    # For bonds, prepend the PAD bond type (using rdkit's BondType.ZERO)
+    bond_labels = sorted(bond_labels)
+    bond_labels = [Chem.rdchem.BondType.ZERO] + bond_labels
+
+    # Create encoder and decoder dictionaries
+    atom_encoder = {l: i for i, l in enumerate(atom_labels)}
+    atom_decoder = {i: l for i, l in enumerate(atom_labels)}
+    bond_encoder = {l: i for i, l in enumerate(bond_labels)}
+    bond_decoder = {i: l for i, l in enumerate(bond_labels)}
+
+    # Ensure directories exist
+    os.makedirs(enc_dir, exist_ok=True)
+    os.makedirs(dec_dir, exist_ok=True)
+
+    # Save the encoders/decoders to disk
+    with open(atom_encoder_path, "wb") as f:
+        pickle.dump(atom_encoder, f)
+    with open(atom_decoder_path, "wb") as f:
+        pickle.dump(atom_decoder, f)
+    with open(bond_encoder_path, "wb") as f:
+        pickle.dump(bond_encoder, f)
+    with open(bond_decoder_path, "wb") as f:
+        pickle.dump(bond_decoder, f)
+
+    print("Encoders/decoders created and saved.")
+    return atom_encoder, atom_decoder, bond_encoder, bond_decoder
+
+def load_molecules(data=None, b_dim=32, m_dim=32, device=None, batch_size=32):
+    data = data.to(device)
+    a = geoutils.to_dense_adj(
+        edge_index = data.edge_index,
+        batch=data.batch,
+        edge_attr=data.edge_attr,
+        max_num_nodes=int(data.batch.shape[0]/batch_size)
+    )
+    x_tensor = data.x.view(batch_size,int(data.batch.shape[0]/batch_size),-1)
+    a_tensor = label2onehot(a, b_dim, device)
+
+    a_tensor_vec = a_tensor.reshape(batch_size,-1)
+    x_tensor_vec = x_tensor.reshape(batch_size,-1)
+    real_graphs = torch.concat((x_tensor_vec,a_tensor_vec),dim=-1)
+
+    return real_graphs, a_tensor, x_tensor

src/model/layers.py

@@ -0,0 +1,234 @@
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class MLP(nn.Module):
+    """
+    A simple Multi-Layer Perceptron (MLP) module consisting of two linear layers with a ReLU activation in between,
+    followed by a dropout on the output.
+
+    Attributes:
+        fc1 (nn.Linear): The first fully-connected layer.
+        act (nn.ReLU): ReLU activation function.
+        fc2 (nn.Linear): The second fully-connected layer.
+        droprateout (nn.Dropout): Dropout layer applied to the output.
+    """
+    def __init__(self, in_feat, hid_feat=None, out_feat=None, dropout=0.):
+        """
+        Initializes the MLP module.
+
+        Args:
+            in_feat (int): Number of input features.
+            hid_feat (int, optional): Number of hidden features. Defaults to in_feat if not provided.
+            out_feat (int, optional): Number of output features. Defaults to in_feat if not provided.
+            dropout (float, optional): Dropout rate. Defaults to 0.
+        """
+        super().__init__()
+
+        # Set hidden and output dimensions to input dimension if not specified
+        if not hid_feat:
+            hid_feat = in_feat
+        if not out_feat:
+            out_feat = in_feat
+
+        self.fc1 = nn.Linear(in_feat, hid_feat)
+        self.act = nn.ReLU()
+        self.fc2 = nn.Linear(hid_feat, out_feat)
+        self.droprateout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        """
+        Forward pass for the MLP.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor after applying the linear layers, activation, and dropout.
+        """
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return self.droprateout(x)
+
+class MHA(nn.Module):
+    """
+    Multi-Head Attention (MHA) module of the graph transformer with edge features incorporated into the attention computation.
+
+    Attributes:
+        heads (int): Number of attention heads.
+        scale (float): Scaling factor for the attention scores.
+        q, k, v (nn.Linear): Linear layers to project the node features into query, key, and value embeddings.
+        e (nn.Linear): Linear layer to project the edge features.
+        d_k (int): Dimension of each attention head.
+        out_e (nn.Linear): Linear layer applied to the computed edge features.
+        out_n (nn.Linear): Linear layer applied to the aggregated node features.
+    """
+    def __init__(self, dim, heads, attention_dropout=0.):
+        """
+        Initializes the Multi-Head Attention module.
+
+        Args:
+            dim (int): Dimensionality of the input features.
+            heads (int): Number of attention heads.
+            attention_dropout (float, optional): Dropout rate for attention (not used explicitly in this implementation).
+        """
+        super().__init__()
+
+        # Ensure that dimension is divisible by the number of heads
+        assert dim % heads == 0
+
+        self.heads = heads
+        self.scale = 1. / math.sqrt(dim)  # Scaling factor for attention
+        # Linear layers for projecting node features
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        # Linear layer for projecting edge features
+        self.e = nn.Linear(dim, dim)
+        self.d_k = dim // heads  # Dimension per head
+
+        # Linear layers for output transformations
+        self.out_e = nn.Linear(dim, dim)
+        self.out_n = nn.Linear(dim, dim)
+
+    def forward(self, node, edge):
+        """
+        Forward pass for the Multi-Head Attention.
+
+        Args:
+            node (torch.Tensor): Node feature tensor of shape (batch, num_nodes, dim).
+            edge (torch.Tensor): Edge feature tensor of shape (batch, num_nodes, num_nodes, dim).
+
+        Returns:
+            tuple: (updated node features, updated edge features)
+        """
+        b, n, c = node.shape
+
+        # Compute query, key, and value embeddings and reshape for multi-head attention
+        q_embed = self.q(node).view(b, n, self.heads, c // self.heads)
+        k_embed = self.k(node).view(b, n, self.heads, c // self.heads)
+        v_embed = self.v(node).view(b, n, self.heads, c // self.heads)
+
+        # Compute edge embeddings
+        e_embed = self.e(edge).view(b, n, n, self.heads, c // self.heads)
+
+        # Adjust dimensions for broadcasting: add singleton dimensions to queries and keys
+        q_embed = q_embed.unsqueeze(2)  # Shape: (b, n, 1, heads, c//heads)
+        k_embed = k_embed.unsqueeze(1)  # Shape: (b, 1, n, heads, c//heads)
+
+        # Compute  attention scores
+        attn = q_embed * k_embed
+        attn = attn / math.sqrt(self.d_k)
+        attn = attn * (e_embed + 1) * e_embed   # Modulated attention incorporating edge features
+
+        edge_out = self.out_e(attn.flatten(3))  # Flatten last dimension for linear layer
+
+        # Apply softmax over the node dimension to obtain normalized attention weights
+        attn = F.softmax(attn, dim=2)
+
+        v_embed = v_embed.unsqueeze(1)  # Adjust dimensions to broadcast: (b, 1, n, heads, c//heads)
+        v_embed = attn * v_embed
+        v_embed = v_embed.sum(dim=2).flatten(2)
+        node_out = self.out_n(v_embed)
+
+        return node_out, edge_out
+
+class Encoder_Block(nn.Module):
+    """
+    Transformer encoder block that integrates node and edge features.
+    
+    Consists of:
+        - A multi-head attention layer with edge modulation.
+        - Two MLP layers, each with residual connections and layer normalization.
+    
+    Attributes:
+        ln1, ln3, ln4, ln5, ln6 (nn.LayerNorm): Layer normalization modules.
+        attn (MHA): Multi-head attention module.
+        mlp, mlp2 (MLP): MLP modules for further transformation of node and edge features.
+    """
+    def __init__(self, dim, heads, act, mlp_ratio=4, drop_rate=0.):
+        """
+        Initializes the encoder block.
+
+        Args:
+            dim (int): Dimensionality of the input features.
+            heads (int): Number of attention heads.
+            act (callable): Activation function (not explicitly used in this block, but provided for potential extensions).
+            mlp_ratio (int, optional): Ratio to determine the hidden layer size in the MLP. Defaults to 4.
+            drop_rate (float, optional): Dropout rate applied in the MLPs. Defaults to 0.
+        """
+        super().__init__()
+
+        self.ln1 = nn.LayerNorm(dim)
+        self.attn = MHA(dim, heads, drop_rate)
+        self.ln3 = nn.LayerNorm(dim)
+        self.ln4 = nn.LayerNorm(dim)
+        self.mlp = MLP(dim, dim * mlp_ratio, dim, dropout=drop_rate)
+        self.mlp2 = MLP(dim, dim * mlp_ratio, dim, dropout=drop_rate)
+        self.ln5 = nn.LayerNorm(dim)
+        self.ln6 = nn.LayerNorm(dim)
+
+    def forward(self, x, y):
+        """
+        Forward pass of the encoder block.
+
+        Args:
+            x (torch.Tensor): Node feature tensor.
+            y (torch.Tensor): Edge feature tensor.
+
+        Returns:
+            tuple: (updated node features, updated edge features)
+        """
+        x1 = self.ln1(x)
+        x2, y1 = self.attn(x1, y)
+        x2 = x1 + x2
+        y2 = y + y1
+        x2 = self.ln3(x2)
+        y2 = self.ln4(y2)
+        x = self.ln5(x2 + self.mlp(x2))
+        y = self.ln6(y2 + self.mlp2(y2))
+        return x, y
+
+class TransformerEncoder(nn.Module):
+    """
+    Transformer Encoder composed of a sequence of encoder blocks.
+
+    Attributes:
+        Encoder_Blocks (nn.ModuleList): A list of Encoder_Block modules stacked sequentially.
+    """
+    def __init__(self, dim, depth, heads, act, mlp_ratio=4, drop_rate=0.1):
+        """
+        Initializes the Transformer Encoder.
+
+        Args:
+            dim (int): Dimensionality of the input features.
+            depth (int): Number of encoder blocks to stack.
+            heads (int): Number of attention heads in each block.
+            act (callable): Activation function (passed to encoder blocks for potential use).
+            mlp_ratio (int, optional): Ratio for determining the hidden layer size in MLP modules. Defaults to 4.
+            drop_rate (float, optional): Dropout rate for the MLPs within each block. Defaults to 0.1.
+        """
+        super().__init__()
+
+        self.Encoder_Blocks = nn.ModuleList([
+            Encoder_Block(dim, heads, act, mlp_ratio, drop_rate)
+            for _ in range(depth)
+        ])
+
+    def forward(self, x, y):
+        """
+        Forward pass of the Transformer Encoder.
+
+        Args:
+            x (torch.Tensor): Node feature tensor.
+            y (torch.Tensor): Edge feature tensor.
+
+        Returns:
+            tuple: (final node features, final edge features) after processing through all encoder blocks.
+        """
+        for block in self.Encoder_Blocks:
+            x, y = block(x, y)
+        return x, y

src/model/loss.py

@@ -0,0 +1,85 @@
+import torch
+
+
+def gradient_penalty(discriminator, real_node, real_edge, fake_node, fake_edge, batch_size, device):
+    """
+    Calculate gradient penalty for WGAN-GP.
+    
+    Args:
+        discriminator: The discriminator model
+        real_node: Real node features
+        real_edge: Real edge features
+        fake_node: Generated node features
+        fake_edge: Generated edge features
+        batch_size: Batch size
+        device: Device to compute on
+        
+    Returns:
+        Gradient penalty term
+    """
+    # Generate random interpolation factors
+    eps_edge = torch.rand(batch_size, 1, 1, 1, device=device)
+    eps_node = torch.rand(batch_size, 1, 1, device=device)
+    
+    # Create interpolated samples
+    int_node = (eps_node * real_node + (1 - eps_node) * fake_node).requires_grad_(True)
+    int_edge = (eps_edge * real_edge + (1 - eps_edge) * fake_edge).requires_grad_(True)
+
+    logits_interpolated = discriminator(int_edge, int_node)
+
+    # Calculate gradients for both node and edge inputs
+    weight = torch.ones(logits_interpolated.size(), requires_grad=False).to(device)
+    gradients = torch.autograd.grad(
+        outputs=logits_interpolated,
+        inputs=[int_node, int_edge],
+        grad_outputs=weight,
+        create_graph=True,
+        retain_graph=True,
+        only_inputs=True
+    )
+
+    # Combine gradients from both inputs
+    gradients_node = gradients[0].view(batch_size, -1)
+    gradients_edge = gradients[1].view(batch_size, -1)
+    gradients = torch.cat([gradients_node, gradients_edge], dim=1)
+
+    # Calculate gradient penalty
+    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
+    
+    return gradient_penalty
+
+
+def discriminator_loss(generator, discriminator, drug_adj, drug_annot, mol_adj, mol_annot, batch_size, device, lambda_gp):
+    # Compute loss for drugs
+    logits_real_disc = discriminator(drug_adj, drug_annot)
+
+    # Use mean reduction for more stable training
+    prediction_real = -torch.mean(logits_real_disc)
+
+    # Compute loss for generated molecules
+    node, edge, node_sample, edge_sample = generator(mol_adj, mol_annot)
+
+    logits_fake_disc = discriminator(edge_sample.detach(), node_sample.detach())
+
+    prediction_fake = torch.mean(logits_fake_disc)
+
+    # Compute gradient penalty using the new function
+    gp = gradient_penalty(discriminator, drug_annot, drug_adj, node_sample.detach(), edge_sample.detach(), batch_size, device)
+
+    # Calculate total discriminator loss
+    d_loss = prediction_fake + prediction_real + lambda_gp * gp
+
+    return node, edge, d_loss
+
+
+def generator_loss(generator, discriminator, mol_adj, mol_annot, batch_size):
+    # Generate fake molecules
+    node, edge, node_sample, edge_sample = generator(mol_adj, mol_annot)
+
+    # Compute logits for fake molecules
+    logits_fake_disc = discriminator(edge_sample, node_sample)
+
+    prediction_fake = -torch.mean(logits_fake_disc)
+    g_loss = prediction_fake
+
+    return g_loss, node, edge, node_sample, edge_sample

src/model/models.py

@@ -0,0 +1,269 @@
+import torch
+import torch.nn as nn
+from src.model.layers import TransformerEncoder
+
+class Generator(nn.Module):
+    """
+    Generator network that uses a Transformer Encoder to process node and edge features.
+    
+    The network first processes input node and edge features with separate linear layers,
+    then applies a Transformer Encoder to model interactions, and finally outputs both transformed
+    features and readout samples.
+    """
+    def __init__(self, act, vertexes, edges, nodes, dropout, dim, depth, heads, mlp_ratio):
+        """
+        Initializes the Generator.
+
+        Args:
+            act (str): Type of activation function to use ("relu", "leaky", "sigmoid", or "tanh").
+            vertexes (int): Number of vertexes in the graph.
+            edges (int): Number of edge features.
+            nodes (int): Number of node features.
+            dropout (float): Dropout rate.
+            dim (int): Dimensionality used for intermediate features.
+            depth (int): Number of Transformer encoder blocks.
+            heads (int): Number of attention heads in the Transformer.
+            mlp_ratio (int): Ratio for determining hidden layer size in MLP modules.
+        """
+        super(Generator, self).__init__()
+        self.vertexes = vertexes
+        self.edges = edges
+        self.nodes = nodes
+        self.depth = depth
+        self.dim = dim
+        self.heads = heads
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+
+        # Set the activation function based on the provided string
+        if act == "relu":
+            act = nn.ReLU()
+        elif act == "leaky":
+            act = nn.LeakyReLU()
+        elif act == "sigmoid":
+            act = nn.Sigmoid()
+        elif act == "tanh":
+            act = nn.Tanh()
+
+        # Calculate the total number of features and dimensions for transformer
+        self.features = vertexes * vertexes * edges + vertexes * nodes
+        self.transformer_dim = vertexes * vertexes * dim + vertexes * dim
+
+        self.node_layers = nn.Sequential(
+            nn.Linear(nodes, 64), act,
+            nn.Linear(64, dim), act,
+            nn.Dropout(self.dropout)
+        )
+        self.edge_layers = nn.Sequential(
+            nn.Linear(edges, 64), act,
+            nn.Linear(64, dim), act,
+            nn.Dropout(self.dropout)
+        )
+        self.TransformerEncoder = TransformerEncoder(
+            dim=self.dim, depth=self.depth, heads=self.heads, act=act,
+            mlp_ratio=self.mlp_ratio, drop_rate=self.dropout
+        )
+
+        self.readout_e = nn.Linear(self.dim, edges)
+        self.readout_n = nn.Linear(self.dim, nodes)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, z_e, z_n):
+        """
+        Forward pass of the Generator.
+        
+        Args:
+            z_e (torch.Tensor): Edge features tensor of shape (batch, vertexes, vertexes, edges).
+            z_n (torch.Tensor): Node features tensor of shape (batch, vertexes, nodes).
+        
+        Returns:
+            tuple: A tuple containing:
+                - node: Updated node features after the transformer.
+                - edge: Updated edge features after the transformer.
+                - node_sample: Readout sample from node features.
+                - edge_sample: Readout sample from edge features.
+        """
+        b, n, c = z_n.shape
+        # The fourth dimension of edge features
+        _, _, _, d = z_e.shape
+
+        # Process node and edge features through their respective layers
+        node = self.node_layers(z_n)
+        edge = self.edge_layers(z_e)
+        # Symmetrize the edge features by averaging with its transpose along vertex dimensions
+        edge = (edge + edge.permute(0, 2, 1, 3)) / 2
+
+        # Pass the features through the Transformer Encoder
+        node, edge = self.TransformerEncoder(node, edge)
+
+        # Readout layers to generate final outputs
+        node_sample = self.readout_n(node)
+        edge_sample = self.readout_e(edge)
+
+        return node, edge, node_sample, edge_sample
+
+
+class Discriminator(nn.Module):
+    """
+    Discriminator network that evaluates node and edge features.
+    
+    It processes features with linear layers, applies a Transformer Encoder to capture dependencies,
+    and finally predicts a scalar value using an MLP on aggregated node features.
+
+    This class is used in DrugGEN model.
+    """
+    def __init__(self, act, vertexes, edges, nodes, dropout, dim, depth, heads, mlp_ratio):
+        """
+        Initializes the Discriminator.
+
+        Args:
+            act (str): Activation function type ("relu", "leaky", "sigmoid", or "tanh").
+            vertexes (int): Number of vertexes.
+            edges (int): Number of edge features.
+            nodes (int): Number of node features.
+            dropout (float): Dropout rate.
+            dim (int): Dimensionality for intermediate representations.
+            depth (int): Number of Transformer encoder blocks.
+            heads (int): Number of attention heads.
+            mlp_ratio (int): MLP ratio for hidden layer dimensions.
+        """
+        super(Discriminator, self).__init__()
+        self.vertexes = vertexes
+        self.edges = edges
+        self.nodes = nodes
+        self.depth = depth
+        self.dim = dim
+        self.heads = heads
+        self.mlp_ratio = mlp_ratio
+        self.dropout = dropout
+
+        # Set the activation function
+        if act == "relu":
+            act = nn.ReLU()
+        elif act == "leaky":
+            act = nn.LeakyReLU()
+        elif act == "sigmoid":
+            act = nn.Sigmoid()
+        elif act == "tanh":
+            act = nn.Tanh()
+
+        self.features = vertexes * vertexes * edges + vertexes * nodes
+        self.transformer_dim = vertexes * vertexes * dim + vertexes * dim
+
+        # Define layers for processing node and edge features
+        self.node_layers = nn.Sequential(
+            nn.Linear(nodes, 64), act,
+            nn.Linear(64, dim), act,
+            nn.Dropout(self.dropout)
+        )
+        self.edge_layers = nn.Sequential(
+            nn.Linear(edges, 64), act,
+            nn.Linear(64, dim), act,
+            nn.Dropout(self.dropout)
+        )
+        # Transformer Encoder for modeling node and edge interactions
+        self.TransformerEncoder = TransformerEncoder(
+            dim=self.dim, depth=self.depth, heads=self.heads, act=act,
+            mlp_ratio=self.mlp_ratio, drop_rate=self.dropout
+        )
+        # Calculate dimensions for node features aggregation
+        self.node_features = vertexes * dim
+        self.edge_features = vertexes * vertexes * dim
+        # MLP to predict a scalar value from aggregated node features
+        self.node_mlp = nn.Sequential(
+            nn.Linear(self.node_features, 64), act,
+            nn.Linear(64, 32), act,
+            nn.Linear(32, 16), act,
+            nn.Linear(16, 1)
+        )
+
+    def forward(self, z_e, z_n):
+        """
+        Forward pass of the Discriminator.
+        
+        Args:
+            z_e (torch.Tensor): Edge features tensor of shape (batch, vertexes, vertexes, edges).
+            z_n (torch.Tensor): Node features tensor of shape (batch, vertexes, nodes).
+        
+        Returns:
+            torch.Tensor: Prediction scores (typically a scalar per sample).
+        """
+        b, n, c = z_n.shape
+        # Unpack the shape of edge features (not used further directly)
+        _, _, _, d = z_e.shape
+
+        # Process node and edge features separately
+        node = self.node_layers(z_n)
+        edge = self.edge_layers(z_e)
+        # Symmetrize edge features by averaging with its transpose
+        edge = (edge + edge.permute(0, 2, 1, 3)) / 2
+
+        # Process features through the Transformer Encoder
+        node, edge = self.TransformerEncoder(node, edge)
+
+        # Flatten node features for MLP
+        node = node.view(b, -1)
+        # Predict a scalar score using the node MLP
+        prediction = self.node_mlp(node)
+
+        return prediction
+
+
+class simple_disc(nn.Module):
+    """
+    A simplified discriminator that processes flattened features through an MLP
+    to predict a scalar score.
+
+    This class is used in NoTarget model.
+    """
+    def __init__(self, act, m_dim, vertexes, b_dim):
+        """
+        Initializes the simple discriminator.
+
+        Args:
+            act (str): Activation function type ("relu", "leaky", "sigmoid", or "tanh").
+            m_dim (int): Dimensionality for atom type features.
+            vertexes (int): Number of vertexes.
+            b_dim (int): Dimensionality for bond type features.
+        """
+        super().__init__()
+
+        # Set the activation function and check if it's supported
+        if act == "relu":
+            act = nn.ReLU()
+        elif act == "leaky":
+            act = nn.LeakyReLU()
+        elif act == "sigmoid":
+            act = nn.Sigmoid()
+        elif act == "tanh":
+            act = nn.Tanh()
+        else:
+            raise ValueError("Unsupported activation function: {}".format(act))
+
+        # Compute total number of features combining both dimensions
+        features = vertexes * m_dim + vertexes * vertexes * b_dim
+        print(vertexes)
+        print(m_dim)
+        print(b_dim)
+        print(features)
+        self.predictor = nn.Sequential(
+            nn.Linear(features, 256), act,
+            nn.Linear(256, 128), act,
+            nn.Linear(128, 64), act,
+            nn.Linear(64, 32), act,
+            nn.Linear(32, 16), act,
+            nn.Linear(16, 1)
+        )
+
+    def forward(self, x):
+        """
+        Forward pass of the simple discriminator.
+        
+        Args:
+            x (torch.Tensor): Input features tensor.
+        
+        Returns:
+            torch.Tensor: Prediction scores.
+        """
+        prediction = self.predictor(x)
+        return prediction

src/util/smiles_cor.py

@@ -0,0 +1,1284 @@
+import os
+import time
+import random
+import re
+import itertools
+import statistics
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchtext.data import TabularDataset, Field, BucketIterator, Iterator
+
+from rdkit import Chem, rdBase, RDLogger
+from rdkit.Chem import (
+    MolStandardize,
+    GraphDescriptors,
+    Lipinski,
+    AllChem,
+)
+from rdkit.Chem.rdSLNParse import MolFromSLN
+from rdkit.Chem.rdmolfiles import MolFromSmiles
+from chembl_structure_pipeline import standardizer
+
+RDLogger.DisableLog('rdApp.*')
+
+SEED = 42
+random.seed(SEED)
+torch.manual_seed(SEED)
+torch.backends.cudnn.deterministic = True
+
+##################################################################################################
+##################################################################################################
+#                                                                                                #
+#  THIS SCRIPT IS DIRECTLY ADAPTED FROM https://github.com/LindeSchoenmaker/SMILES-corrector     #
+#                                                                                                #
+##################################################################################################
+##################################################################################################
+def is_smiles(array,
+              TRG,
+              reverse: bool,
+              return_output=False,
+              src=None,
+              src_field=None):
+    """Turns predicted tokens within batch into smiles and evaluates their validity
+    Arguments:
+        array: Tensor with most probable token for each location for each sequence in batch
+            [trg len, batch size]
+        TRG: target field for getting tokens from vocab
+        reverse (bool): True if the target sequence is reversed
+        return_output (bool): True if output sequences and their validity should be saved
+    Returns:
+        df: dataframe with correct and incorrect sequences
+        valids: list with booleans that show if prediction was a valid SMILES (True) or invalid one (False)
+        smiless: list of the predicted smiles
+    """
+    trg_field = TRG
+    valids = []
+    smiless = []
+    if return_output:
+        df = pd.DataFrame()
+    else:
+        df = None
+    batch_size = array.size(1)
+    # check if the first token should be removed, first token is zero because
+    # outputs initaliazed to all be zeros
+    if int((array[0, 0]).tolist()) == 0:
+        start = 1
+    else:
+        start = 0
+    # for each sequence in the batch
+    for i in range(0, batch_size):
+        # turns sequence from tensor to list skipps first row as this is not
+        # filled in in forward
+        sequence = (array[start:, i]).tolist()
+        # goes from embedded to tokens
+        trg_tokens = [trg_field.vocab.itos[int(t)] for t in sequence]
+        # print(trg_tokens)
+        # takes all tokens untill eos token, model would be faster if did this
+        # one step earlier, but then changes in vocab order would disrupt.
+        rev_tokens = list(
+            itertools.takewhile(lambda x: x != "<eos>", trg_tokens))
+        if reverse:
+            rev_tokens = rev_tokens[::-1]
+        smiles = "".join(rev_tokens)
+        # determine how many valid smiles are made
+        valid = True if MolFromSmiles(smiles) else False
+        valids.append(valid)
+        smiless.append(smiles)
+        if return_output:
+            if valid:
+                df.loc[i, "CORRECT"] = smiles
+            else:
+                df.loc[i, "INCORRECT"] = smiles
+
+    # add the original drugex outputs to the _de dataframe
+    if return_output and src is not None:
+        for i in range(0, batch_size):
+            # turns sequence from tensor to list skipps first row as this is
+            # <sos> for src
+            sequence = (src[1:, i]).tolist()
+            # goes from embedded to tokens
+            src_tokens = [src_field.vocab.itos[int(t)] for t in sequence]
+            # takes all tokens untill eos token, model would be faster if did
+            # this one step earlier, but then changes in vocab order would
+            # disrupt.
+            rev_tokens = list(
+                itertools.takewhile(lambda x: x != "<eos>", src_tokens))
+            smiles = "".join(rev_tokens)
+            df.loc[i, "ORIGINAL"] = smiles
+
+    return df, valids, smiless
+
+
+def is_unchanged(array,
+                 TRG,
+                 reverse: bool,
+                 return_output=False,
+                 src=None,
+                 src_field=None):
+    """Checks is output is different from input
+    Arguments:
+        array: Tensor with most probable token for each location for each sequence in batch
+            [trg len, batch size]
+        TRG: target field for getting tokens from vocab
+        reverse (bool): True if the target sequence is reversed
+        return_output (bool): True if output sequences and their validity should be saved
+    Returns:
+        df: dataframe with correct and incorrect sequences
+        valids: list with booleans that show if prediction was a valid SMILES (True) or invalid one (False)
+        smiless: list of the predicted smiles
+    """
+    trg_field = TRG
+    sources = []
+    batch_size = array.size(1)
+    unchanged = 0
+
+    # check if the first token should be removed, first token is zero because
+    # outputs initaliazed to all be zeros
+    if int((array[0, 0]).tolist()) == 0:
+        start = 1
+    else:
+        start = 0
+
+    for i in range(0, batch_size):
+        # turns sequence from tensor to list skipps first row as this is <sos>
+        # for src
+        sequence = (src[1:, i]).tolist()
+        # goes from embedded to tokens
+        src_tokens = [src_field.vocab.itos[int(t)] for t in sequence]
+        # takes all tokens untill eos token, model would be faster if did this
+        # one step earlier, but then changes in vocab order would disrupt.
+        rev_tokens = list(
+            itertools.takewhile(lambda x: x != "<eos>", src_tokens))
+        smiles = "".join(rev_tokens)
+        sources.append(smiles)
+
+    # for each sequence in the batch
+    for i in range(0, batch_size):
+        # turns sequence from tensor to list skipps first row as this is not
+        # filled in in forward
+        sequence = (array[start:, i]).tolist()
+        # goes from embedded to tokens
+        trg_tokens = [trg_field.vocab.itos[int(t)] for t in sequence]
+        # print(trg_tokens)
+        # takes all tokens untill eos token, model would be faster if did this
+        # one step earlier, but then changes in vocab order would disrupt.
+        rev_tokens = list(
+            itertools.takewhile(lambda x: x != "<eos>", trg_tokens))
+        if reverse:
+            rev_tokens = rev_tokens[::-1]
+        smiles = "".join(rev_tokens)
+        # determine how many valid smiles are made
+        valid = True if MolFromSmiles(smiles) else False
+        if not valid:
+            if smiles == sources[i]:
+                unchanged += 1
+
+    return unchanged
+
+
+def molecule_reconstruction(array, TRG, reverse: bool, outputs):
+    """Turns target tokens within batch into smiles and compares them to predicted output smiles
+    Arguments:
+        array: Tensor with target's token for each location for each sequence in batch
+            [trg len, batch size]
+        TRG: target field for getting tokens from vocab
+        reverse (bool): True if the target sequence is reversed
+        outputs: list of predicted SMILES sequences
+    Returns:
+         matches(int): number of total right molecules
+    """
+    trg_field = TRG
+    matches = 0
+    targets = []
+    batch_size = array.size(1)
+    # for each sequence in the batch
+    for i in range(0, batch_size):
+        # turns sequence from tensor to list skipps first row as this is not
+        # filled in in forward
+        sequence = (array[1:, i]).tolist()
+        # goes from embedded to tokens
+        trg_tokens = [trg_field.vocab.itos[int(t)] for t in sequence]
+        # takes all tokens untill eos token, model would be faster if did this
+        # one step earlier, but then changes in vocab order would disrupt.
+        rev_tokens = list(
+            itertools.takewhile(lambda x: x != "<eos>", trg_tokens))
+        if reverse:
+            rev_tokens = rev_tokens[::-1]
+        smiles = "".join(rev_tokens)
+        targets.append(smiles)
+    for i in range(0, batch_size):
+        m = MolFromSmiles(targets[i])
+        p = MolFromSmiles(outputs[i])
+        if p is not None:
+            if m.HasSubstructMatch(p) and p.HasSubstructMatch(m):
+                matches += 1
+    return matches
+
+
+def complexity_whitlock(mol: Chem.Mol, includeAllDescs=False):
+    """
+    Complexity as defined in DOI:10.1021/jo9814546
+    S: complexity = 4*#rings + 2*#unsat + #hetatm + 2*#chiral
+    Other descriptors:
+        H: size = #bonds (Hydrogen atoms included)
+        G: S + H
+        Ratio: S / H
+    """
+    mol_ = Chem.Mol(mol)
+    nrings = Lipinski.RingCount(mol_) - Lipinski.NumAromaticRings(mol_)
+    Chem.rdmolops.SetAromaticity(mol_)
+    unsat = sum(1 for bond in mol_.GetBonds()
+                if bond.GetBondTypeAsDouble() == 2)
+    hetatm = len(mol_.GetSubstructMatches(Chem.MolFromSmarts("[!#6]")))
+    AllChem.EmbedMolecule(mol_)
+    Chem.AssignAtomChiralTagsFromStructure(mol_)
+    chiral = len(Chem.FindMolChiralCenters(mol_))
+    S = 4 * nrings + 2 * unsat + hetatm + 2 * chiral
+    if not includeAllDescs:
+        return S
+    Chem.rdmolops.Kekulize(mol_)
+    mol_ = Chem.AddHs(mol_)
+    H = sum(bond.GetBondTypeAsDouble() for bond in mol_.GetBonds())
+    G = S + H
+    R = S / H
+    return {"WhitlockS": S, "WhitlockH": H, "WhitlockG": G, "WhitlockRatio": R}
+
+
+def complexity_baronechanon(mol: Chem.Mol):
+    """
+    Complexity as defined in DOI:10.1021/ci000145p
+    """
+    mol_ = Chem.Mol(mol)
+    Chem.Kekulize(mol_)
+    Chem.RemoveStereochemistry(mol_)
+    mol_ = Chem.RemoveHs(mol_, updateExplicitCount=True)
+    degree, counts = 0, 0
+    for atom in mol_.GetAtoms():
+        degree += 3 * 2**(atom.GetExplicitValence() - atom.GetNumExplicitHs() -
+                          1)
+        counts += 3 if atom.GetSymbol() == "C" else 6
+    ringterm = sum(map(lambda x: 6 * len(x), mol_.GetRingInfo().AtomRings()))
+    return degree + counts + ringterm
+
+
+def calc_complexity(array,
+                    TRG,
+                    reverse,
+                    valids,
+                    complexity_function=GraphDescriptors.BertzCT):
+    """Calculates the complexity of inputs that are not correct.
+    Arguments:
+        array: Tensor with target's token for each location for each sequence in batch
+            [trg len, batch size]
+        TRG: target field for getting tokens from vocab
+        reverse (bool): True if the target sequence is reversed
+        valids: list with booleans that show if prediction was a valid SMILES (True) or invalid one (False)
+        complexity_function: the type of complexity measure that will be used
+            GraphDescriptors.BertzCT
+            complexity_whitlock
+            complexity_baronechanon
+    Returns:
+         matches(int): mean of complexity values
+    """
+    trg_field = TRG
+    sources = []
+    complexities = []
+    loc = torch.BoolTensor(valids)
+    # only keeps rows in batch size dimension where valid is false
+    array = array[:, loc == False]
+    # should check if this still works
+    # array = torch.transpose(array, 0, 1)
+    array_size = array.size(1)
+    for i in range(0, array_size):
+        # turns sequence from tensor to list skipps first row as this is not
+        # filled in in forward
+        sequence = (array[1:, i]).tolist()
+        # goes from embedded to tokens
+        trg_tokens = [trg_field.vocab.itos[int(t)] for t in sequence]
+        # takes all tokens untill eos token, model would be faster if did this
+        # one step earlier, but then changes in vocab order would disrupt.
+        rev_tokens = list(
+            itertools.takewhile(lambda x: x != "<eos>", trg_tokens))
+        if reverse:
+            rev_tokens = rev_tokens[::-1]
+        smiles = "".join(rev_tokens)
+        sources.append(smiles)
+    for source in sources:
+        try:
+            m = MolFromSmiles(source)
+        except BaseException:
+            m = MolFromSLN(source)
+        complexities.append(complexity_function(m))
+    if len(complexities) > 0:
+        mean = statistics.mean(complexities)
+    else:
+        mean = 0
+    return mean
+
+
+def epoch_time(start_time, end_time):
+    elapsed_time = end_time - start_time
+    elapsed_mins = int(elapsed_time / 60)
+    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
+    return elapsed_mins, elapsed_secs
+
+
+class Convo:
+    """Class for training and evaluating transformer and convolutional neural network
+    
+    Methods
+    -------
+    train_model()
+        train model for initialized number of epochs
+    evaluate(return_output)
+        use model with validation loader (& optionally drugex loader) to get test loss & other metrics
+    translate(loader)
+        translate inputs from loader (different from evaluate in that no target sequence is used)
+    """
+
+    def train_model(self):
+        optimizer = optim.Adam(self.parameters(), lr=self.lr)
+        log = open(f"{self.out}.log", "a")
+        best_error = np.inf
+        for epoch in range(self.epochs):
+            self.train()
+            start_time = time.time()
+            loss_train = 0
+            for i, batch in enumerate(self.loader_train):
+                optimizer.zero_grad()
+                # changed src,trg call to match with bentrevett
+                # src, trg = batch['src'], batch['trg']
+                trg = batch.trg
+                src = batch.src
+                output, attention = self(src, trg[:, :-1])
+                # feed the source and target into def forward to get the output
+                # Xuhan uses forward for this, with istrain = true
+                output_dim = output.shape[-1]
+                # changed
+                output = output.contiguous().view(-1, output_dim)
+                trg = trg[:, 1:].contiguous().view(-1)
+                # output = output[:,:,0]#.view(-1)
+                # output = output[1:].view(-1, output.shape[-1])
+                # trg = trg[1:].view(-1)
+                loss = nn.CrossEntropyLoss(
+                    ignore_index=self.TRG.vocab.stoi[self.TRG.pad_token])
+                a, b = output.view(-1), trg.to(self.device).view(-1)
+                # changed
+                # loss = loss(output.view(0), trg.view(0).to(device))
+                loss = loss(output, trg)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.parameters(), self.clip)
+                optimizer.step()
+                loss_train += loss.item()
+                # turned off for now, as not using voc so won't work, output is a tensor
+                # output = [(trg len - 1) * batch size, output dim]
+                # smiles, valid = is_valid_smiles(output, reversed)
+                # if valid:
+                #    valids += 1
+                #    smiless.append(smiles)
+            # added .dataset becaue len(iterator) gives len(self.dataset) /
+            # self.batch_size)
+            loss_train /= len(self.loader_train)
+            info = f"Epoch: {epoch+1:02} step: {i} loss_train: {loss_train:.4g}"
+            # model is used to generate trg based on src from the validation set to assess performance
+            # similar to Xuhan, although he doesn't use the if loop
+            if self.loader_valid is not None:
+                return_output = False
+                if epoch + 1 == self.epochs:
+                    return_output = True
+                (
+                    valids,
+                    loss_valid,
+                    valids_de,
+                    df_output,
+                    df_output_de,
+                    right_molecules,
+                    complexity,
+                    unchanged,
+                    unchanged_de,
+                ) = self.evaluate(return_output)
+                reconstruction_error = 1 - right_molecules / len(
+                    self.loader_valid.dataset)
+                error = 1 - valids / len(self.loader_valid.dataset)
+                complexity = complexity / len(self.loader_valid)
+                unchan = unchanged / (len(self.loader_valid.dataset) - valids)
+                info += f" loss_valid: {loss_valid:.4g} error_rate: {error:.4g} molecule_reconstruction_error_rate: {reconstruction_error:.4g} unchanged: {unchan:.4g} invalid_target_complexity: {complexity:.4g}"
+                if self.loader_drugex is not None:
+                    error_de = 1 - valids_de / len(self.loader_drugex.dataset)
+                    unchan_de = unchanged_de / (
+                        len(self.loader_drugex.dataset) - valids_de)
+                    info += f" error_rate_drugex: {error_de:.4g} unchanged_drugex: {unchan_de:.4g}"
+
+                if reconstruction_error < best_error:
+                    torch.save(self.state_dict(), f"{self.out}.pkg")
+                    best_error = reconstruction_error
+                    last_save = epoch
+                else:
+                    if epoch - last_save >= 10 and best_error != 1:
+                        torch.save(self.state_dict(), f"{self.out}_last.pkg")
+                        (
+                            valids,
+                            loss_valid,
+                            valids_de,
+                            df_output,
+                            df_output_de,
+                            right_molecules,
+                            complexity,
+                            unchanged,
+                            unchanged_de,
+                        ) = self.evaluate(True)
+                        end_time = time.time()
+                        epoch_mins, epoch_secs = epoch_time(
+                            start_time, end_time)
+                        info += f" Time: {epoch_mins}m {epoch_secs}s"
+                 
+                        break
+            elif error < best_error:
+                torch.save(self.state_dict(), f"{self.out}.pkg")
+                best_error = error
+            end_time = time.time()
+            epoch_mins, epoch_secs = epoch_time(start_time, end_time)
+            info += f" Time: {epoch_mins}m {epoch_secs}s"
+   
+        
+        torch.save(self.state_dict(), f"{self.out}_last.pkg")
+        log.close()
+        self.load_state_dict(torch.load(f"{self.out}.pkg"))
+        df_output.to_csv(f"{self.out}.csv", index=False)
+        df_output_de.to_csv(f"{self.out}_de.csv", index=False)
+
+    def evaluate(self, return_output):
+        self.eval()
+        test_loss = 0
+        df_output = pd.DataFrame()
+        df_output_de = pd.DataFrame()
+        valids = 0
+        valids_de = 0
+        unchanged = 0
+        unchanged_de = 0
+        right_molecules = 0
+        complexity = 0
+        with torch.no_grad():
+            for _, batch in enumerate(self.loader_valid):
+                trg = batch.trg
+                src = batch.src
+                output, attention = self.forward(src, trg[:, :-1])
+                pred_token = output.argmax(2)
+                array = torch.transpose(pred_token, 0, 1)
+                trg_trans = torch.transpose(trg, 0, 1)
+                output_dim = output.shape[-1]
+                output = output.contiguous().view(-1, output_dim)
+                trg = trg[:, 1:].contiguous().view(-1)
+                src_trans = torch.transpose(src, 0, 1)
+                df_batch, valid, smiless = is_smiles(
+                    array, self.TRG, reverse=True, return_output=return_output)
+                unchanged += is_unchanged(
+                    array,
+                    self.TRG,
+                    reverse=True,
+                    return_output=return_output,
+                    src=src_trans,
+                    src_field=self.SRC,
+                )
+                matches = molecule_reconstruction(trg_trans,
+                                                  self.TRG,
+                                                  reverse=True,
+                                                  outputs=smiless)
+                complexity += calc_complexity(trg_trans,
+                                              self.TRG,
+                                              reverse=True,
+                                              valids=valid)
+                if df_batch is not None:
+                    df_output = pd.concat([df_output, df_batch],
+                                          ignore_index=True)
+                right_molecules += matches
+                valids += sum(valid)
+                # trg = trg[1:].view(-1)
+                # output, trg = output[1:].view(-1, output.shape[-1]), trg[1:].view(-1)
+                loss = nn.CrossEntropyLoss(
+                    ignore_index=self.TRG.vocab.stoi[self.TRG.pad_token])
+                loss = loss(output, trg)
+            test_loss += loss.item()
+            if self.loader_drugex is not None:
+                for _, batch in enumerate(self.loader_drugex):
+                    src = batch.src
+                    output = self.translate_sentence(src, self.TRG,
+                                                     self.device)
+                    # checks the number of valid smiles
+                    pred_token = output.argmax(2)
+                    array = torch.transpose(pred_token, 0, 1)
+                    src_trans = torch.transpose(src, 0, 1)
+                    df_batch, valid, smiless = is_smiles(
+                        array,
+                        self.TRG,
+                        reverse=True,
+                        return_output=return_output,
+                        src=src_trans,
+                        src_field=self.SRC,
+                    )
+                    unchanged_de += is_unchanged(
+                        array,
+                        self.TRG,
+                        reverse=True,
+                        return_output=return_output,
+                        src=src_trans,
+                        src_field=self.SRC,
+                    )
+                    if df_batch is not None:
+                        df_output_de = pd.concat([df_output_de, df_batch],
+                                                 ignore_index=True)
+                    valids_de += sum(valid)
+        return (
+            valids,
+            test_loss / len(self.loader_valid),
+            valids_de,
+            df_output,
+            df_output_de,
+            right_molecules,
+            complexity,
+            unchanged,
+            unchanged_de,
+        )
+
+    def translate(self, loader):
+        self.eval()
+        df_output_de = pd.DataFrame()
+        valids_de = 0
+        with torch.no_grad():
+            for _, batch in enumerate(loader):
+                src = batch.src
+                output = self.translate_sentence(src, self.TRG, self.device)
+                # checks the number of valid smiles
+                pred_token = output.argmax(2)
+                array = torch.transpose(pred_token, 0, 1)
+                src_trans = torch.transpose(src, 0, 1)
+                df_batch, valid, smiless = is_smiles(
+                    array,
+                    self.TRG,
+                    reverse=True,
+                    return_output=True,
+                    src=src_trans,
+                    src_field=self.SRC,
+                )
+                if df_batch is not None:
+                    df_output_de = pd.concat([df_output_de, df_batch],
+                                             ignore_index=True)
+                valids_de += sum(valid)
+        return valids_de, df_output_de
+
+
+class Encoder(nn.Module):
+
+    def __init__(self, input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout,
+                 max_length, device):
+        super().__init__()
+        self.device = device
+        self.tok_embedding = nn.Embedding(input_dim, hid_dim)
+        self.pos_embedding = nn.Embedding(max_length, hid_dim)
+        self.layers = nn.ModuleList([
+            EncoderLayer(hid_dim, n_heads, pf_dim, dropout, device)
+            for _ in range(n_layers)
+        ])
+
+        self.dropout = nn.Dropout(dropout)
+        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)
+
+    def forward(self, src, src_mask):
+        # src = [batch size, src len]
+        # src_mask = [batch size, src len]
+        batch_size = src.shape[0]
+        src_len = src.shape[1]
+        pos = (torch.arange(0, src_len).unsqueeze(0).repeat(batch_size,
+                                                            1).to(self.device))
+        # pos = [batch size, src len]
+        src = self.dropout((self.tok_embedding(src) * self.scale) +
+                           self.pos_embedding(pos))
+        # src = [batch size, src len, hid dim]
+        for layer in self.layers:
+            src = layer(src, src_mask)
+        # src = [batch size, src len, hid dim]
+        return src
+
+
+class EncoderLayer(nn.Module):
+
+    def __init__(self, hid_dim, n_heads, pf_dim, dropout, device):
+        super().__init__()
+
+        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)
+        self.ff_layer_norm = nn.LayerNorm(hid_dim)
+        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads,
+                                                      dropout, device)
+        self.positionwise_feedforward = PositionwiseFeedforwardLayer(
+            hid_dim, pf_dim, dropout)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, src, src_mask):
+        # src = [batch size, src len, hid dim]
+        # src_mask = [batch size, src len]
+        # self attention
+        _src, _ = self.self_attention(src, src, src, src_mask)
+        # dropout, residual connection and layer norm
+        src = self.self_attn_layer_norm(src + self.dropout(_src))
+        # src = [batch size, src len, hid dim]
+        # positionwise feedforward
+        _src = self.positionwise_feedforward(src)
+        # dropout, residual and layer norm
+        src = self.ff_layer_norm(src + self.dropout(_src))
+        # src = [batch size, src len, hid dim]
+
+        return src
+
+
+class MultiHeadAttentionLayer(nn.Module):
+
+    def __init__(self, hid_dim, n_heads, dropout, device):
+        super().__init__()
+        assert hid_dim % n_heads == 0
+        self.hid_dim = hid_dim
+        self.n_heads = n_heads
+        self.head_dim = hid_dim // n_heads
+        self.fc_q = nn.Linear(hid_dim, hid_dim)
+        self.fc_k = nn.Linear(hid_dim, hid_dim)
+        self.fc_v = nn.Linear(hid_dim, hid_dim)
+        self.fc_o = nn.Linear(hid_dim, hid_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)
+
+    def forward(self, query, key, value, mask=None):
+        batch_size = query.shape[0]
+        # query = [batch size, query len, hid dim]
+        # key = [batch size, key len, hid dim]
+        # value = [batch size, value len, hid dim]
+        Q = self.fc_q(query)
+        K = self.fc_k(key)
+        V = self.fc_v(value)
+        # Q = [batch size, query len, hid dim]
+        # K = [batch size, key len, hid dim]
+        # V = [batch size, value len, hid dim]
+        Q = Q.view(batch_size, -1, self.n_heads,
+                   self.head_dim).permute(0, 2, 1, 3)
+        K = K.view(batch_size, -1, self.n_heads,
+                   self.head_dim).permute(0, 2, 1, 3)
+        V = V.view(batch_size, -1, self.n_heads,
+                   self.head_dim).permute(0, 2, 1, 3)
+        # Q = [batch size, n heads, query len, head dim]
+        # K = [batch size, n heads, key len, head dim]
+        # V = [batch size, n heads, value len, head dim]
+        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale
+        # energy = [batch size, n heads, query len, key len]
+        if mask is not None:
+            energy = energy.masked_fill(mask == 0, -1e10)
+        attention = torch.softmax(energy, dim=-1)
+        # attention = [batch size, n heads, query len, key len]
+        x = torch.matmul(self.dropout(attention), V)
+        # x = [batch size, n heads, query len, head dim]
+        x = x.permute(0, 2, 1, 3).contiguous()
+        # x = [batch size, query len, n heads, head dim]
+        x = x.view(batch_size, -1, self.hid_dim)
+        # x = [batch size, query len, hid dim]
+        x = self.fc_o(x)
+        # x = [batch size, query len, hid dim]
+        return x, attention
+
+
+class PositionwiseFeedforwardLayer(nn.Module):
+
+    def __init__(self, hid_dim, pf_dim, dropout):
+        super().__init__()
+        self.fc_1 = nn.Linear(hid_dim, pf_dim)
+        self.fc_2 = nn.Linear(pf_dim, hid_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # x = [batch size, seq len, hid dim]
+        x = self.dropout(torch.relu(self.fc_1(x)))
+        # x = [batch size, seq len, pf dim]
+        x = self.fc_2(x)
+        # x = [batch size, seq len, hid dim]
+
+        return x
+
+
+class Decoder(nn.Module):
+
+    def __init__(
+        self,
+        output_dim,
+        hid_dim,
+        n_layers,
+        n_heads,
+        pf_dim,
+        dropout,
+        max_length,
+        device,
+    ):
+        super().__init__()
+        self.device = device
+        self.tok_embedding = nn.Embedding(output_dim, hid_dim)
+        self.pos_embedding = nn.Embedding(max_length, hid_dim)
+        self.layers = nn.ModuleList([
+            DecoderLayer(hid_dim, n_heads, pf_dim, dropout, device)
+            for _ in range(n_layers)
+        ])
+        self.fc_out = nn.Linear(hid_dim, output_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)
+
+    def forward(self, trg, enc_src, trg_mask, src_mask):
+        # trg = [batch size, trg len]
+        # enc_src = [batch size, src len, hid dim]
+        # trg_mask = [batch size, trg len]
+        # src_mask = [batch size, src len]
+        batch_size = trg.shape[0]
+        trg_len = trg.shape[1]
+        pos = (torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size,
+                                                            1).to(self.device))
+        # pos = [batch size, trg len]
+        trg = self.dropout((self.tok_embedding(trg) * self.scale) +
+                           self.pos_embedding(pos))
+        # trg = [batch size, trg len, hid dim]
+        for layer in self.layers:
+            trg, attention = layer(trg, enc_src, trg_mask, src_mask)
+        # trg = [batch size, trg len, hid dim]
+        # attention = [batch size, n heads, trg len, src len]
+        output = self.fc_out(trg)
+        # output = [batch size, trg len, output dim]
+        return output, attention
+
+
+class DecoderLayer(nn.Module):
+
+    def __init__(self, hid_dim, n_heads, pf_dim, dropout, device):
+        super().__init__()
+        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)
+        self.enc_attn_layer_norm = nn.LayerNorm(hid_dim)
+        self.ff_layer_norm = nn.LayerNorm(hid_dim)
+        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads,
+                                                      dropout, device)
+        self.encoder_attention = MultiHeadAttentionLayer(
+            hid_dim, n_heads, dropout, device)
+        self.positionwise_feedforward = PositionwiseFeedforwardLayer(
+            hid_dim, pf_dim, dropout)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, trg, enc_src, trg_mask, src_mask):
+        # trg = [batch size, trg len, hid dim]
+        # enc_src = [batch size, src len, hid dim]
+        # trg_mask = [batch size, trg len]
+        # src_mask = [batch size, src len]
+        # self attention
+        _trg, _ = self.self_attention(trg, trg, trg, trg_mask)
+        # dropout, residual connection and layer norm
+        trg = self.self_attn_layer_norm(trg + self.dropout(_trg))
+        # trg = [batch size, trg len, hid dim]
+        # encoder attention
+        _trg, attention = self.encoder_attention(trg, enc_src, enc_src,
+                                                 src_mask)
+        # dropout, residual connection and layer norm
+        trg = self.enc_attn_layer_norm(trg + self.dropout(_trg))
+        # trg = [batch size, trg len, hid dim]
+        # positionwise feedforward
+        _trg = self.positionwise_feedforward(trg)
+        # dropout, residual and layer norm
+        trg = self.ff_layer_norm(trg + self.dropout(_trg))
+        # trg = [batch size, trg len, hid dim]
+        # attention = [batch size, n heads, trg len, src len]
+        return trg, attention
+
+
+class Seq2Seq(nn.Module, Convo):
+
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        src_pad_idx,
+        trg_pad_idx,
+        device,
+        loader_train: DataLoader,
+        out: str,
+        loader_valid=None,
+        loader_drugex=None,
+        epochs=100,
+        lr=0.0005,
+        clip=0.1,
+        reverse=True,
+        TRG=None,
+        SRC=None,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.src_pad_idx = src_pad_idx
+        self.trg_pad_idx = trg_pad_idx
+        self.device = device
+        self.loader_train = loader_train
+        self.out = out
+        self.loader_valid = loader_valid
+        self.loader_drugex = loader_drugex
+        self.epochs = epochs
+        self.lr = lr
+        self.clip = clip
+        self.reverse = reverse
+        self.TRG = TRG
+        self.SRC = SRC
+
+    def make_src_mask(self, src):
+        # src = [batch size, src len]
+        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)
+        # src_mask = [batch size, 1, 1, src len]
+        return src_mask
+
+    def make_trg_mask(self, trg):
+        # trg = [batch size, trg len]
+        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)
+        # trg_pad_mask = [batch size, 1, 1, trg len]
+        trg_len = trg.shape[1]
+        trg_sub_mask = torch.tril(
+            torch.ones((trg_len, trg_len), device=self.device)).bool()
+        # trg_sub_mask = [trg len, trg len]
+        trg_mask = trg_pad_mask & trg_sub_mask
+        # trg_mask = [batch size, 1, trg len, trg len]
+        return trg_mask
+
+    def forward(self, src, trg):
+        # src = [batch size, src len]
+        # trg = [batch size, trg len]
+        src_mask = self.make_src_mask(src)
+        trg_mask = self.make_trg_mask(trg)
+        # src_mask = [batch size, 1, 1, src len]
+        # trg_mask = [batch size, 1, trg len, trg len]
+        enc_src = self.encoder(src, src_mask)
+        # enc_src = [batch size, src len, hid dim]
+        output, attention = self.decoder(trg, enc_src, trg_mask, src_mask)
+        # output = [batch size, trg len, output dim]
+        # attention = [batch size, n heads, trg len, src len]
+        return output, attention
+
+    def translate_sentence(self, src, trg_field, device, max_len=202):
+        self.eval()
+        src_mask = self.make_src_mask(src)
+        with torch.no_grad():
+            enc_src = self.encoder(src, src_mask)
+        trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]
+        batch_size = src.shape[0]
+        trg = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)
+        trg = trg.repeat(batch_size, 1)
+        for i in range(max_len):
+            # turned model into self.
+            trg_mask = self.make_trg_mask(trg)
+            with torch.no_grad():
+                output, attention = self.decoder(trg, enc_src, trg_mask,
+                                                 src_mask)
+            pred_tokens = output.argmax(2)[:, -1].unsqueeze(1)
+            trg = torch.cat((trg, pred_tokens), 1)
+
+        return output
+
+
+def remove_floats(df: pd.DataFrame, subset: str):
+    """Preprocessing step to remove any entries that are not strings"""
+    df_subset = df[subset]
+    df[subset] = df[subset].astype(str)
+    # only keep entries that stayed the same after applying astype str
+    df = df[df[subset] == df_subset].copy()
+
+    return df
+
+
+def smi_tokenizer(smi: str, reverse=False) -> list:
+    """
+    Tokenize a SMILES molecule
+    """
+    pattern = r"(\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9])"
+    regex = re.compile(pattern)
+    # tokens = ['<sos>'] + [token for token in regex.findall(smi)] + ['<eos>']
+    tokens = [token for token in regex.findall(smi)]
+    # assert smi == ''.join(tokens[1:-1])
+    assert smi == "".join(tokens[:])
+    # try:
+    #     assert smi == "".join(tokens[:])
+    # except:
+    #     print(smi)
+    #     print("".join(tokens[:]))
+    if reverse:
+        return tokens[::-1]
+    return tokens
+
+
+def init_weights(m: nn.Module):
+    if hasattr(m, "weight") and m.weight.dim() > 1:
+        nn.init.xavier_uniform_(m.weight.data)
+
+
+def count_parameters(model: nn.Module):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def epoch_time(start_time, end_time):
+    elapsed_time = end_time - start_time
+    elapsed_mins = int(elapsed_time / 60)
+    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
+    return elapsed_mins, elapsed_secs
+
+
+def initialize_model(folder_out: str,
+                     data_source: str,
+                     error_source: str,
+                     device: torch.device,
+                     threshold: int,
+                     epochs: int,
+                     layers: int = 3,
+                     batch_size: int = 16,
+                     invalid_type: str = "all",
+                     num_errors: int = 1,
+                     validation_step=False):
+    """Create encoder decoder models for specified model (currently only translator) & type of invalid SMILES
+
+    param data: collection of invalid, valid SMILES pairs
+    param invalid_smiles_path: path to previously generated invalid SMILES
+    param invalid_type: type of errors introduced into invalid SMILES
+
+    return:
+
+    """
+
+    # set fields
+    SRC = Field(
+        tokenize=lambda x: smi_tokenizer(x),
+        init_token="<sos>",
+        eos_token="<eos>",
+        batch_first=True,
+    )
+    TRG = Field(
+        tokenize=lambda x: smi_tokenizer(x, reverse=True),
+        init_token="<sos>",
+        eos_token="<eos>",
+        batch_first=True,
+    )
+
+    if validation_step:
+        train, val = TabularDataset.splits(
+            path=f'{folder_out}errors/split/',
+            train=f"{data_source}_{invalid_type}_{num_errors}_errors_train.csv",
+            validation=
+            f"{data_source}_{invalid_type}_{num_errors}_errors_dev.csv",
+            format="CSV",
+            skip_header=False,
+            fields={
+                "ERROR": ("src", SRC),
+                "STD_SMILES": ("trg", TRG)
+            },
+        )
+        SRC.build_vocab(train, val, max_size=1000)
+        TRG.build_vocab(train, val, max_size=1000)
+    else:
+        train = TabularDataset(
+            path=
+            f'{folder_out}{data_source}_{invalid_type}_{num_errors}_errors.csv',
+            format="CSV",
+            skip_header=False,
+            fields={
+                "ERROR": ("src", SRC),
+                "STD_SMILES": ("trg", TRG)
+            },
+        )
+        SRC.build_vocab(train, max_size=1000)
+        TRG.build_vocab(train, max_size=1000)
+
+    drugex = TabularDataset(
+        path=error_source,
+        format="csv",
+        skip_header=False,
+        fields={
+            "SMILES": ("src", SRC),
+            "SMILES_TARGET": ("trg", TRG)
+        },
+    )
+
+
+    #SRC.vocab = torch.load('vocab_src.pth')
+    #TRG.vocab = torch.load('vocab_trg.pth')
+
+    # model parameters
+    EPOCHS = epochs
+    BATCH_SIZE = batch_size
+    INPUT_DIM = len(SRC.vocab)
+    OUTPUT_DIM = len(TRG.vocab)
+    HID_DIM = 256
+    ENC_LAYERS = layers
+    DEC_LAYERS = layers
+    ENC_HEADS = 8
+    DEC_HEADS = 8
+    ENC_PF_DIM = 512
+    DEC_PF_DIM = 512
+    ENC_DROPOUT = 0.1
+    DEC_DROPOUT = 0.1
+    SRC_PAD_IDX = SRC.vocab.stoi[SRC.pad_token]
+    TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]
+    # add 2 to length for start and stop tokens
+    MAX_LENGTH = threshold + 2
+
+    # model name
+    MODEL_OUT_FOLDER = f"{folder_out}"
+
+    MODEL_NAME = "transformer_%s_%s_%s_%s_%s" % (
+        invalid_type, num_errors, data_source, BATCH_SIZE, layers)
+    if not os.path.exists(MODEL_OUT_FOLDER):
+        os.mkdir(MODEL_OUT_FOLDER)
+
+    out = os.path.join(MODEL_OUT_FOLDER, MODEL_NAME)
+
+    torch.save(SRC.vocab, f'{out}_vocab_src.pth')
+    torch.save(TRG.vocab, f'{out}_vocab_trg.pth')
+
+    # iterator is a dataloader
+    # iterator to pass to the same length and create batches in which the
+    # amount of padding is minimized
+    if validation_step:
+        train_iter, val_iter = BucketIterator.splits(
+            (train, val),
+            batch_sizes=(BATCH_SIZE, 256),
+            sort_within_batch=True,
+            shuffle=True,
+            # the BucketIterator needs to be told what function it should use to
+            # group the data.
+            sort_key=lambda x: len(x.src),
+            device=device,
+        )
+    else:
+        train_iter = BucketIterator(
+            train,
+            batch_size=BATCH_SIZE,
+            sort_within_batch=True,
+            shuffle=True,
+            # the BucketIterator needs to be told what function it should use to
+            # group the data.
+            sort_key=lambda x: len(x.src),
+            device=device,
+        )
+        val_iter = None
+
+    drugex_iter = Iterator(
+        drugex,
+        batch_size=64,
+        device=device,
+        sort=False,
+        sort_within_batch=True,
+        sort_key=lambda x: len(x.src),
+        repeat=False,
+    )
+
+
+    # model initialization
+
+    enc = Encoder(
+        INPUT_DIM,
+        HID_DIM,
+        ENC_LAYERS,
+        ENC_HEADS,
+        ENC_PF_DIM,
+        ENC_DROPOUT,
+        MAX_LENGTH,
+        device,
+    )
+    dec = Decoder(
+        OUTPUT_DIM,
+        HID_DIM,
+        DEC_LAYERS,
+        DEC_HEADS,
+        DEC_PF_DIM,
+        DEC_DROPOUT,
+        MAX_LENGTH,
+        device,
+    )
+
+    model = Seq2Seq(
+        enc,
+        dec,
+        SRC_PAD_IDX,
+        TRG_PAD_IDX,
+        device,
+        train_iter,
+        out=out,
+        loader_valid=val_iter,
+        loader_drugex=drugex_iter,
+        epochs=EPOCHS,
+        TRG=TRG,
+        SRC=SRC,
+    ).to(device)
+
+
+
+
+    return model, out, SRC
+
+
+def train_model(model, out, assess):
+    """Apply given weights (& assess performance or train further) or start training new model
+
+    Args:
+        model: initialized model
+        out: .pkg file with model parameters
+        asses: bool 
+
+    Returns:
+        model with (new) weights
+    """
+
+    if os.path.exists(f"{out}.pkg") and assess:
+
+
+        model.load_state_dict(torch.load(f=out + ".pkg"))
+        (
+            valids,
+            loss_valid,
+            valids_de,
+            df_output,
+            df_output_de,
+            right_molecules,
+            complexity,
+            unchanged,
+            unchanged_de,
+        ) = model.evaluate(True)
+
+
+        # log = open('unchanged.log', 'a')
+        # info = f'type: comb unchanged: {unchan:.4g} unchanged_drugex: {unchan_de:.4g}'
+        # print(info, file=log, flush = True)
+        # print(valids_de)
+        # print(unchanged_de)
+
+        # print(unchan)
+        # print(unchan_de)
+        # df_output_de.to_csv(f'{out}_de_new.csv', index = False)
+
+        # error_de = 1 - valids_de / len(drugex_iter.dataset)
+        # print(error_de)
+        # df_output.to_csv(f'{out}_par.csv', index = False)
+
+    elif os.path.exists(f"{out}.pkg"):
+
+        # starts from the model after the last epoch, not the best epoch
+        model.load_state_dict(torch.load(f=out + "_last.pkg"))
+        # need to change how log file names epochs
+        model.train_model()
+    else:
+
+        model = model.apply(init_weights)
+        model.train_model()
+
+    return model
+
+
+def correct_SMILES(model, out, error_source, device, SRC):
+    """Model that is given corrects SMILES and return number of correct ouputs and dataframe containing all outputs
+    Args:
+        model: initialized model
+        out: .pkg file with model parameters
+        asses: bool 
+
+    Returns:
+        valids: number of fixed outputs
+        df_output: dataframe containing output (either correct or incorrect) & original input
+    """
+    ## account for tokens that are not yet in SRC without changing existing SRC token embeddings
+    errors = TabularDataset(
+        path=error_source,
+        format="csv",
+        skip_header=False,
+        fields={"SMILES": ("src", SRC)},
+    )
+
+    errors_loader = Iterator(
+        errors,
+        batch_size=64,
+        device=device,
+        sort=False,
+        sort_within_batch=True,
+        sort_key=lambda x: len(x.src),
+        repeat=False,
+    )
+    model.load_state_dict(torch.load(f=out + ".pkg",map_location=torch.device('cpu')))
+    # add option to use different iterator maybe?
+
+    valids, df_output = model.translate(errors_loader)
+    #df_output.to_csv(f"{error_source}_fixed.csv", index=False)
+
+
+    return valids, df_output
+
+
+
+class smi_correct(object):
+    def __init__(self, model_name, trans_file_path):
+    # set random seed, used for error generation & initiation transformer
+    
+        self.SEED = 42
+        random.seed(self.SEED)
+        self.model_name = model_name
+        self.folder_out = "data/"
+        
+        self.trans_file_path = trans_file_path
+
+        if not os.path.exists(self.folder_out):
+            os.makedirs(self.folder_out)
+            
+        self.invalid_type = 'multiple'
+        self.num_errors = 12
+        self.threshold = 200
+        self.data_source = f"PAPYRUS_{self.threshold}"
+        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+        self.initialize_source = 'data/papyrus_rnn_S.csv' # change this path
+        
+    def standardization_pipeline(self, smile):
+        desalter = MolStandardize.rdMolStandardize.LargestFragmentChooser()
+        std_smile = None
+        if not isinstance(smile, str): return None
+        m = Chem.MolFromSmiles(smile)
+        # skips smiles for which no mol file could be generated
+        if m is not None:
+            # standardizes
+            std_m = standardizer.standardize_mol(m)
+            # strips salts
+            std_m_p, exclude = standardizer.get_parent_mol(std_m)
+            if not exclude:
+                # choose largest fragment for rare cases where chembl structure
+                # pipeline leaves 2 fragments
+                std_m_p_d = desalter.choose(std_m_p)
+                std_smile = Chem.MolToSmiles(std_m_p_d)
+        return std_smile      
+    
+    def remove_smiles_duplicates(self, dataframe: pd.DataFrame,
+                             subset: str) -> pd.DataFrame:
+        return dataframe.drop_duplicates(subset=subset)  
+    
+    def correct(self, smi):
+    
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        model, out, SRC = initialize_model(self.folder_out,
+                                        self.data_source,
+                                        error_source=self.initialize_source,
+                                        device=device,
+                                        threshold=self.threshold,
+                                        epochs=30,
+                                        layers=3,
+                                        batch_size=16,
+                                        invalid_type=self.invalid_type,
+                                        num_errors=self.num_errors)
+
+        valids, df_output = correct_SMILES(model, out, smi, device,
+                                        SRC)
+        
+        df_output["SMILES"] = df_output.apply(lambda row: self.standardization_pipeline(row["CORRECT"]), axis=1)
+        
+        df_output = self.remove_smiles_duplicates(df_output, subset="SMILES").drop(columns=["CORRECT", "INCORRECT", "ORIGINAL"]).dropna()
+        
+        return df_output

src/util/utils.py

@@ -0,0 +1,930 @@
+import os
+import time
+import math
+import datetime
+import warnings
+import itertools
+from copy import deepcopy
+from functools import partial
+from collections import Counter
+from multiprocessing import Pool
+from statistics import mean
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.lines import Line2D
+from scipy.spatial.distance import cosine as cos_distance
+
+import torch
+import wandb
+
+from rdkit import Chem, DataStructs, RDLogger
+from rdkit.Chem import (
+    AllChem,
+    Draw,
+    Descriptors,
+    Lipinski,
+    Crippen,
+    rdMolDescriptors,
+    FilterCatalog,
+)
+from rdkit.Chem.Scaffolds import MurckoScaffold
+
+# Disable RDKit warnings
+RDLogger.DisableLog("rdApp.*")
+
+
+class Metrics(object):
+    """
+    Collection of static methods to compute various metrics for molecules.
+    """
+
+    @staticmethod
+    def valid(x):
+        """
+        Checks whether the molecule is valid.
+        
+        Args:
+            x: RDKit molecule object.
+        
+        Returns:
+            bool: True if molecule is valid and has a non-empty SMILES representation.
+        """
+        return x is not None and Chem.MolToSmiles(x) != ''
+
+    @staticmethod
+    def tanimoto_sim_1v2(data1, data2):
+        """
+        Computes the average Tanimoto similarity for paired fingerprints.
+        
+        Args:
+            data1: Fingerprint data for first set.
+            data2: Fingerprint data for second set.
+        
+        Returns:
+            float: The average Tanimoto similarity between corresponding fingerprints.
+        """
+        # Determine the minimum size between two arrays for pairing
+        min_len = data1.size if data1.size > data2.size else data2
+        sims = []
+        for i in range(min_len):
+            sim = DataStructs.FingerprintSimilarity(data1[i], data2[i])
+            sims.append(sim)
+        # Use 'mean' from statistics; note that variable 'sim' was used, corrected to use sims list.
+        mean_sim = mean(sims)
+        return mean_sim
+
+    @staticmethod
+    def mol_length(x):
+        """
+        Computes the length of the largest fragment (by character count) in a SMILES string.
+        
+        Args:
+            x (str): SMILES string.
+        
+        Returns:
+            int: Number of alphabetic characters in the longest fragment of the SMILES.
+        """
+        if x is not None:
+            # Split at dots (.) and take the fragment with maximum length, then count alphabetic characters.
+            return len([char for char in max(x.split(sep="."), key=len).upper() if char.isalpha()])
+        else:
+            return 0
+
+    @staticmethod
+    def max_component(data, max_len):
+        """
+        Returns the average normalized length of molecules in the dataset.
+        
+        Each molecule's length is computed and divided by max_len, then averaged.
+        
+        Args:
+            data (iterable): Collection of SMILES strings.
+            max_len (int): Maximum possible length for normalization.
+        
+        Returns:
+            float: Normalized average length.
+        """
+        lengths = np.array(list(map(Metrics.mol_length, data)), dtype=np.float32)
+        return (lengths / max_len).mean()
+
+    @staticmethod
+    def mean_atom_type(data):
+        """
+        Computes the average number of unique atom types in the provided node data.
+        
+        Args:
+            data (iterable): Iterable containing node data with unique atom types.
+        
+        Returns:
+            float: The average count of unique atom types, subtracting one.
+        """
+        atom_types_used = []
+        for i in data:
+            # Assuming each element i has a .unique() method that returns unique atom types.
+            atom_types_used.append(len(i.unique().tolist()))
+        av_type = np.mean(atom_types_used) - 1
+        return av_type
+
+
+def mols2grid_image(mols, path):
+    """
+    Saves grid images for a list of molecules.
+    
+    For each molecule in the list, computes 2D coordinates and saves an image file.
+    
+    Args:
+        mols (list): List of RDKit molecule objects.
+        path (str): Directory where images will be saved.
+    """
+    # Replace None molecules with an empty molecule
+    mols = [e if e is not None else Chem.RWMol() for e in mols]
+
+    for i in range(len(mols)):
+        if Metrics.valid(mols[i]):
+            AllChem.Compute2DCoords(mols[i])
+            file_path = os.path.join(path, "{}.png".format(i + 1))
+            Draw.MolToFile(mols[i], file_path, size=(1200, 1200))
+            # wandb.save(file_path)  # Optionally save to Weights & Biases
+        else:
+            continue
+
+
+def save_smiles_matrices(mols, edges_hard, nodes_hard, path, data_source=None):
+    """
+    Saves the edge and node matrices along with SMILES strings to text files.
+    
+    Each file contains the edge matrix, node matrix, and SMILES representation for a molecule.
+    
+    Args:
+        mols (list): List of RDKit molecule objects.
+        edges_hard (torch.Tensor): Tensor of edge features.
+        nodes_hard (torch.Tensor): Tensor of node features.
+        path (str): Directory where files will be saved.
+        data_source: Optional data source information (not used in function).
+    """
+    mols = [e if e is not None else Chem.RWMol() for e in mols]
+
+    for i in range(len(mols)):
+        if Metrics.valid(mols[i]):
+            save_path = os.path.join(path, "{}.txt".format(i + 1))
+            with open(save_path, "a") as f:
+                np.savetxt(f, edges_hard[i].cpu().numpy(), header="edge matrix:\n", fmt='%1.2f')
+                f.write("\n")
+                np.savetxt(f, nodes_hard[i].cpu().numpy(), header="node matrix:\n", footer="\nsmiles:", fmt='%1.2f')
+                f.write("\n")
+            # Append the SMILES representation to the file
+            with open(save_path, "a") as f:
+                print(Chem.MolToSmiles(mols[i]), file=f)
+            # wandb.save(save_path)  # Optionally save to Weights & Biases
+        else:
+            continue
+
+def dense_to_sparse_with_attr(adj):
+    """
+    Converts a dense adjacency matrix to a sparse representation.
+    
+    Args:
+        adj (torch.Tensor): Adjacency matrix tensor (2D or 3D) with square last two dimensions.
+    
+    Returns:
+        tuple: A tuple containing indices and corresponding edge attributes.
+    """
+    assert adj.dim() >= 2 and adj.dim() <= 3
+    assert adj.size(-1) == adj.size(-2)
+
+    index = adj.nonzero(as_tuple=True)
+    edge_attr = adj[index]
+
+    if len(index) == 3:
+        batch = index[0] * adj.size(-1)
+        index = (batch + index[1], batch + index[2])
+    return index, edge_attr
+
+
+def mol_sample(sample_directory, edges, nodes, idx, i, matrices2mol, dataset_name):
+    """
+    Samples molecules from edge and node predictions, then saves grid images and text files.
+    
+    Args:
+        sample_directory (str): Directory to save the samples.
+        edges (torch.Tensor): Edge predictions tensor.
+        nodes (torch.Tensor): Node predictions tensor.
+        idx (int): Current index for naming the sample.
+        i (int): Epoch/iteration index.
+        matrices2mol (callable): Function to convert matrices to RDKit molecule.
+        dataset_name (str): Name of the dataset for file naming.
+    """
+    sample_path = os.path.join(sample_directory, "{}_{}-epoch_iteration".format(idx + 1, i + 1))
+    # Get the index of the maximum predicted feature along the last dimension
+    g_edges_hat_sample = torch.max(edges, -1)[1]
+    g_nodes_hat_sample = torch.max(nodes, -1)[1]
+    # Convert matrices to molecule objects
+    mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(),
+                        strict=True, file_name=dataset_name)
+           for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
+
+    if not os.path.exists(sample_path):
+        os.makedirs(sample_path)
+
+    mols2grid_image(mol, sample_path)
+    save_smiles_matrices(mol, g_edges_hat_sample.detach(), g_nodes_hat_sample.detach(), sample_path)
+
+    # Remove the directory if no files were saved
+    if len(os.listdir(sample_path)) == 0:
+        os.rmdir(sample_path)
+
+    print("Valid molecules are saved.")
+    print("Valid matrices and smiles are saved")
+
+
+def logging(log_path, start_time, i, idx, loss, save_path, drug_smiles, edge, node, 
+            matrices2mol, dataset_name, real_adj, real_annot, drug_vecs):
+    """
+    Logs training statistics and evaluation metrics.
+    
+    The function generates molecules from predictions, computes various metrics such as
+    validity, uniqueness, novelty, and similarity scores, and logs them using wandb and a file.
+    
+    Args:
+        log_path (str): Path to save the log file.
+        start_time (float): Start time to compute elapsed time.
+        i (int): Current iteration index.
+        idx (int): Current epoch index.
+        loss (dict): Dictionary to update with loss and metric values.
+        save_path (str): Directory path to save sample outputs.
+        drug_smiles (list): List of reference drug SMILES.
+        edge (torch.Tensor): Edge prediction tensor.
+        node (torch.Tensor): Node prediction tensor.
+        matrices2mol (callable): Function to convert matrices to molecules.
+        dataset_name (str): Dataset name.
+        real_adj (torch.Tensor): Ground truth adjacency matrix tensor.
+        real_annot (torch.Tensor): Ground truth annotation tensor.
+        drug_vecs (list): List of drug vectors for similarity calculation.
+    """
+    g_edges_hat_sample = torch.max(edge, -1)[1]
+    g_nodes_hat_sample = torch.max(node, -1)[1]
+
+    a_tensor_sample = torch.max(real_adj, -1)[1].float()
+    x_tensor_sample = torch.max(real_annot, -1)[1].float()
+
+    # Generate molecules from predictions and real data
+    mols = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(),
+                         strict=True, file_name=dataset_name)
+            for e_, n_ in zip(g_edges_hat_sample, g_nodes_hat_sample)]
+    real_mol = [matrices2mol(n_.data.cpu().numpy(), e_.data.cpu().numpy(),
+                              strict=True, file_name=dataset_name)
+                for e_, n_ in zip(a_tensor_sample, x_tensor_sample)]
+
+    # Compute average number of atom types
+    atom_types_average = Metrics.mean_atom_type(g_nodes_hat_sample)
+    real_smiles = [Chem.MolToSmiles(x) for x in real_mol if x is not None]
+    gen_smiles = []
+    uniq_smiles = []
+    for line in mols:
+        if line is not None:
+            gen_smiles.append(Chem.MolToSmiles(line))
+            uniq_smiles.append(Chem.MolToSmiles(line))
+        elif line is None:
+            gen_smiles.append(None)
+
+    # Process SMILES to take the longest fragment if multiple are present
+    gen_smiles_saves = [None if x is None else max(x.split('.'), key=len) for x in gen_smiles]
+    uniq_smiles_saves = [None if x is None else max(x.split('.'), key=len) for x in uniq_smiles]
+
+    # Save the generated SMILES to a text file
+    sample_save_dir = os.path.join(save_path, "samples.txt")
+    with open(sample_save_dir, "a") as f:
+        for s in gen_smiles_saves:
+            if s is not None:
+                f.write(s + "\n")
+
+    k = len(set(uniq_smiles_saves) - {None})
+    et = time.time() - start_time
+    et = str(datetime.timedelta(seconds=et))[:-7]
+    log_str = "Elapsed [{}], Epoch/Iteration [{}/{}]".format(et, idx, i + 1)
+    
+    # Generate molecular fingerprints for similarity computations
+    gen_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in mols if x is not None]
+    chembl_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in real_mol if x is not None]
+
+    # Compute evaluation metrics: validity, uniqueness, novelty, similarity scores, and average maximum molecule length.
+    valid = fraction_valid(gen_smiles_saves)
+    unique = fraction_unique(uniq_smiles_saves, k)
+    novel_starting_mol = novelty(gen_smiles_saves, real_smiles)
+    novel_akt = novelty(gen_smiles_saves, drug_smiles)
+    if len(uniq_smiles_saves) == 0:
+        snn_chembl = 0
+        snn_akt = 0
+        maxlen = 0
+    else:
+        snn_chembl = average_agg_tanimoto(np.array(chembl_vecs), np.array(gen_vecs))
+        snn_akt = average_agg_tanimoto(np.array(drug_vecs), np.array(gen_vecs))
+        maxlen = Metrics.max_component(uniq_smiles_saves, 45)
+
+    # Update loss dictionary with computed metrics
+    loss.update({
+        'Validity': valid,
+        'Uniqueness': unique,
+        'Novelty': novel_starting_mol,
+        'Novelty_akt': novel_akt,
+        'SNN_chembl': snn_chembl,
+        'SNN_akt': snn_akt,
+        'MaxLen': maxlen,
+        'Atom_types': atom_types_average
+    })
+
+    # Log metrics using wandb
+    wandb.log({
+        "Validity": valid,
+        "Uniqueness": unique,
+        "Novelty": novel_starting_mol,
+        "Novelty_akt": novel_akt,
+        "SNN_chembl": snn_chembl,
+        "SNN_akt": snn_akt,
+        "MaxLen": maxlen,
+        "Atom_types": atom_types_average
+    })
+
+    # Append each metric to the log string and write to the log file
+    for tag, value in loss.items():
+        log_str += ", {}: {:.4f}".format(tag, value)
+    with open(log_path, "a") as f:
+        f.write(log_str + "\n")
+    print(log_str)
+    print("\n")
+
+
+def plot_grad_flow(named_parameters, model, itera, epoch, grad_flow_directory):
+    """
+    Plots the gradients flowing through different layers during training.
+    
+    This is useful to check for possible gradient vanishing or exploding problems.
+    
+    Args:
+        named_parameters (iterable): Iterable of (name, parameter) tuples from the model.
+        model (str): Name of the model (used for saving the plot).
+        itera (int): Iteration index.
+        epoch (int): Current epoch.
+        grad_flow_directory (str): Directory to save the gradient flow plot.
+    """
+    ave_grads = []
+    max_grads = []
+    layers = []
+    for n, p in named_parameters:
+        if p.requires_grad and ("bias" not in n):
+            layers.append(n)
+            ave_grads.append(p.grad.abs().mean().cpu())
+            max_grads.append(p.grad.abs().max().cpu())
+    # Plot maximum gradients and average gradients for each layer
+    plt.bar(np.arange(len(max_grads)), max_grads, alpha=0.1, lw=1, color="c")
+    plt.bar(np.arange(len(max_grads)), ave_grads, alpha=0.1, lw=1, color="b")
+    plt.hlines(0, 0, len(ave_grads) + 1, lw=2, color="k")
+    plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical")
+    plt.xlim(left=0, right=len(ave_grads))
+    plt.ylim(bottom=-0.001, top=1)  # Zoom in on lower gradient regions
+    plt.xlabel("Layers")
+    plt.ylabel("Average Gradient")
+    plt.title("Gradient Flow")
+    plt.grid(True)
+    plt.legend([
+        Line2D([0], [0], color="c", lw=4),
+        Line2D([0], [0], color="b", lw=4),
+        Line2D([0], [0], color="k", lw=4)
+    ], ['max-gradient', 'mean-gradient', 'zero-gradient'])
+    # Save the plot to the specified directory
+    plt.savefig(os.path.join(grad_flow_directory, "weights_" + model + "_" + str(itera) + "_" + str(epoch) + ".png"), dpi=500, bbox_inches='tight')
+
+
+def get_mol(smiles_or_mol):
+    """
+    Loads a SMILES string or molecule into an RDKit molecule object.
+    
+    Args:
+        smiles_or_mol (str or RDKit Mol): SMILES string or RDKit molecule.
+    
+    Returns:
+        RDKit Mol or None: Sanitized molecule object, or None if invalid.
+    """
+    if isinstance(smiles_or_mol, str):
+        if len(smiles_or_mol) == 0:
+            return None
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+        if mol is None:
+            return None
+        try:
+            Chem.SanitizeMol(mol)
+        except ValueError:
+            return None
+        return mol
+    return smiles_or_mol
+
+
+def mapper(n_jobs):
+    """
+    Returns a mapping function for parallel or serial processing.
+    
+    If n_jobs == 1, returns the built-in map function.
+    If n_jobs > 1, returns a function that uses a multiprocessing pool.
+    
+    Args:
+        n_jobs (int or pool object): Number of jobs or a Pool instance.
+    
+    Returns:
+        callable: A function that acts like map.
+    """
+    if n_jobs == 1:
+        def _mapper(*args, **kwargs):
+            return list(map(*args, **kwargs))
+        return _mapper
+    if isinstance(n_jobs, int):
+        pool = Pool(n_jobs)
+        def _mapper(*args, **kwargs):
+            try:
+                result = pool.map(*args, **kwargs)
+            finally:
+                pool.terminate()
+            return result
+        return _mapper
+    return n_jobs.map
+
+
+def remove_invalid(gen, canonize=True, n_jobs=1):
+    """
+    Removes invalid molecules from the provided dataset.
+    
+    Optionally canonizes the SMILES strings.
+    
+    Args:
+        gen (list): List of SMILES strings.
+        canonize (bool): Whether to convert to canonical SMILES.
+        n_jobs (int): Number of parallel jobs.
+    
+    Returns:
+        list: Filtered list of valid molecules.
+    """
+    if not canonize:
+        mols = mapper(n_jobs)(get_mol, gen)
+        return [gen_ for gen_, mol in zip(gen, mols) if mol is not None]
+    return [x for x in mapper(n_jobs)(canonic_smiles, gen) if x is not None]
+
+
+def fraction_valid(gen, n_jobs=1):
+    """
+    Computes the fraction of valid molecules in the dataset.
+    
+    Args:
+        gen (list): List of SMILES strings.
+        n_jobs (int): Number of parallel jobs.
+    
+    Returns:
+        float: Fraction of molecules that are valid.
+    """
+    gen = mapper(n_jobs)(get_mol, gen)
+    return 1 - gen.count(None) / len(gen)
+
+
+def canonic_smiles(smiles_or_mol):
+    """
+    Converts a SMILES string or molecule to its canonical SMILES.
+    
+    Args:
+        smiles_or_mol (str or RDKit Mol): Input molecule.
+    
+    Returns:
+        str or None: Canonical SMILES string or None if invalid.
+    """
+    mol = get_mol(smiles_or_mol)
+    if mol is None:
+        return None
+    return Chem.MolToSmiles(mol)
+
+
+def fraction_unique(gen, k=None, n_jobs=1, check_validity=True):
+    """
+    Computes the fraction of unique molecules.
+    
+    Optionally computes unique@k, where only the first k molecules are considered.
+    
+    Args:
+        gen (list): List of SMILES strings.
+        k (int): Optional cutoff for unique@k computation.
+        n_jobs (int): Number of parallel jobs.
+        check_validity (bool): Whether to check for validity of molecules.
+    
+    Returns:
+        float: Fraction of unique molecules.
+    """
+    if k is not None:
+        if len(gen) < k:
+            warnings.warn("Can't compute unique@{}.".format(k) +
+                          " gen contains only {} molecules".format(len(gen)))
+        gen = gen[:k]
+    if check_validity:
+        canonic = list(mapper(n_jobs)(canonic_smiles, gen))
+        canonic = [i for i in canonic if i is not None]
+    set_cannonic = set(canonic)
+    return 0 if len(canonic) == 0 else len(set_cannonic) / len(canonic)
+
+
+def novelty(gen, train, n_jobs=1):
+    """
+    Computes the novelty score of generated molecules.
+    
+    Novelty is defined as the fraction of generated molecules that do not appear in the training set.
+    
+    Args:
+        gen (list): List of generated SMILES strings.
+        train (list): List of training SMILES strings.
+        n_jobs (int): Number of parallel jobs.
+    
+    Returns:
+        float: Novelty score.
+    """
+    gen_smiles = mapper(n_jobs)(canonic_smiles, gen)
+    gen_smiles_set = set(gen_smiles) - {None}
+    train_set = set(train)
+    return 0 if len(gen_smiles_set) == 0 else len(gen_smiles_set - train_set) / len(gen_smiles_set)
+
+
+def internal_diversity(gen):
+    """
+    Computes the internal diversity of a set of molecules.
+    
+    Internal diversity is defined as one minus the average Tanimoto similarity between all pairs.
+    
+    Args:
+        gen: Array-like representation of molecules.
+    
+    Returns:
+        tuple: Mean and standard deviation of internal diversity.
+    """
+    diversity = [1 - x for x in average_agg_tanimoto(gen, gen, agg="mean", intdiv=True)]
+    return np.mean(diversity), np.std(diversity)
+
+
+def average_agg_tanimoto(stock_vecs, gen_vecs, batch_size=5000, agg='max', device='cpu', p=1, intdiv=False):
+    """
+    Computes the average aggregated Tanimoto similarity between two sets of molecular fingerprints.
+    
+    For each fingerprint in gen_vecs, finds the closest (max or mean) similarity with fingerprints in stock_vecs.
+    
+    Args:
+        stock_vecs (numpy.ndarray): Array of fingerprint vectors from the reference set.
+        gen_vecs (numpy.ndarray): Array of fingerprint vectors from the generated set.
+        batch_size (int): Batch size for processing fingerprints.
+        agg (str): Aggregation method, either 'max' or 'mean'.
+        device (str): Device to perform computations on.
+        p (int): Power for averaging.
+        intdiv (bool): Whether to return individual similarities or the average.
+    
+    Returns:
+        float or numpy.ndarray: Average aggregated Tanimoto similarity or array of individual scores.
+    """
+    assert agg in ['max', 'mean'], "Can aggregate only max or mean"
+    agg_tanimoto = np.zeros(len(gen_vecs))
+    total = np.zeros(len(gen_vecs))
+    for j in range(0, stock_vecs.shape[0], batch_size):
+        x_stock = torch.tensor(stock_vecs[j:j + batch_size]).to(device).float()
+        for i in range(0, gen_vecs.shape[0], batch_size):
+            y_gen = torch.tensor(gen_vecs[i:i + batch_size]).to(device).float()
+            y_gen = y_gen.transpose(0, 1)
+            tp = torch.mm(x_stock, y_gen)
+            # Compute Jaccard/Tanimoto similarity
+            jac = (tp / (x_stock.sum(1, keepdim=True) + y_gen.sum(0, keepdim=True) - tp)).cpu().numpy()
+            jac[np.isnan(jac)] = 1
+            if p != 1:
+                jac = jac ** p
+            if agg == 'max':
+                agg_tanimoto[i:i + y_gen.shape[1]] = np.maximum(
+                    agg_tanimoto[i:i + y_gen.shape[1]], jac.max(0))
+            elif agg == 'mean':
+                agg_tanimoto[i:i + y_gen.shape[1]] += jac.sum(0)
+                total[i:i + y_gen.shape[1]] += jac.shape[0]
+    if agg == 'mean':
+        agg_tanimoto /= total
+    if p != 1:
+        agg_tanimoto = (agg_tanimoto) ** (1 / p)
+    if intdiv:
+        return agg_tanimoto
+    else:
+        return np.mean(agg_tanimoto)
+
+
+def str2bool(v):
+    """
+    Converts a string to a boolean.
+    
+    Args:
+        v (str): Input string.
+    
+    Returns:
+        bool: True if the string is 'true' (case insensitive), else False.
+    """
+    return v.lower() in ('true')
+
+
+def obey_lipinski(mol):
+    """
+    Checks if a molecule obeys Lipinski's Rule of Five.
+    
+    The function evaluates weight, hydrogen bond donors and acceptors, logP, and rotatable bonds.
+    
+    Args:
+        mol (RDKit Mol): Molecule object.
+    
+    Returns:
+        int: Number of Lipinski rules satisfied.
+    """
+    mol = deepcopy(mol)
+    Chem.SanitizeMol(mol)
+    rule_1 = Descriptors.ExactMolWt(mol) < 500
+    rule_2 = Lipinski.NumHDonors(mol) <= 5
+    rule_3 = Lipinski.NumHAcceptors(mol) <= 10
+    rule_4 = (logp := Crippen.MolLogP(mol) >= -2) & (logp <= 5)
+    rule_5 = Chem.rdMolDescriptors.CalcNumRotatableBonds(mol) <= 10
+    return np.sum([int(a) for a in [rule_1, rule_2, rule_3, rule_4, rule_5]])
+
+
+def obey_veber(mol):
+    """
+    Checks if a molecule obeys Veber's rules.
+    
+    Veber's rules focus on the number of rotatable bonds and topological polar surface area.
+    
+    Args:
+        mol (RDKit Mol): Molecule object.
+    
+    Returns:
+        int: Number of Veber's rules satisfied.
+    """
+    mol = deepcopy(mol)
+    Chem.SanitizeMol(mol)
+    rule_1 = rdMolDescriptors.CalcNumRotatableBonds(mol) <= 10
+    rule_2 = rdMolDescriptors.CalcTPSA(mol) <= 140
+    return np.sum([int(a) for a in [rule_1, rule_2]])
+
+
+def load_pains_filters():
+    """
+    Loads the PAINS (Pan-Assay INterference compoundS) filters A, B, and C.
+    
+    Returns:
+        FilterCatalog: An RDKit FilterCatalog object containing PAINS filters.
+    """
+    params = FilterCatalog.FilterCatalogParams()
+    params.AddCatalog(FilterCatalog.FilterCatalogParams.FilterCatalogs.PAINS_A)
+    params.AddCatalog(FilterCatalog.FilterCatalogParams.FilterCatalogs.PAINS_B)
+    params.AddCatalog(FilterCatalog.FilterCatalogParams.FilterCatalogs.PAINS_C)
+    catalog = FilterCatalog.FilterCatalog(params)
+    return catalog
+
+
+def is_pains(mol, catalog):
+    """
+    Checks if the given molecule is a PAINS compound.
+    
+    Args:
+        mol (RDKit Mol): Molecule object.
+        catalog (FilterCatalog): A catalog of PAINS filters.
+    
+    Returns:
+        bool: True if the molecule matches a PAINS filter, else False.
+    """
+    entry = catalog.GetFirstMatch(mol)
+    return entry is not None
+
+
+def mapper(n_jobs):
+    """
+    Returns a mapping function for parallel or serial processing.
+    
+    If n_jobs == 1, returns the built-in map function.
+    If n_jobs > 1, returns a function that uses a multiprocessing pool.
+    
+    Args:
+        n_jobs (int or pool object): Number of jobs or a Pool instance.
+    
+    Returns:
+        callable: A function that acts like map.
+    """
+    if n_jobs == 1:
+        def _mapper(*args, **kwargs):
+            return list(map(*args, **kwargs))
+        return _mapper
+    if isinstance(n_jobs, int):
+        pool = Pool(n_jobs)
+        def _mapper(*args, **kwargs):
+            try:
+                result = pool.map(*args, **kwargs)
+            finally:
+                pool.terminate()
+            return result
+        return _mapper
+    return n_jobs.map
+
+
+def fragmenter(mol):
+    """
+    Fragments a molecule using BRICS and returns a list of fragment SMILES.
+    
+    Args:
+        mol (str or RDKit Mol): Input molecule.
+    
+    Returns:
+        list: List of fragment SMILES strings.
+    """
+    fgs = AllChem.FragmentOnBRICSBonds(get_mol(mol))
+    fgs_smi = Chem.MolToSmiles(fgs).split(".")
+    return fgs_smi
+
+
+def get_mol(smiles_or_mol):
+    """
+    Loads a SMILES string or molecule into an RDKit molecule object.
+    
+    Args:
+        smiles_or_mol (str or RDKit Mol): SMILES string or molecule.
+    
+    Returns:
+        RDKit Mol or None: Sanitized molecule object or None if invalid.
+    """
+    if isinstance(smiles_or_mol, str):
+        if len(smiles_or_mol) == 0:
+            return None
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+        if mol is None:
+            return None
+        try:
+            Chem.SanitizeMol(mol)
+        except ValueError:
+            return None
+        return mol
+    return smiles_or_mol
+
+
+def compute_fragments(mol_list, n_jobs=1):
+    """
+    Fragments a list of molecules using BRICS and returns a counter of fragment occurrences.
+    
+    Args:
+        mol_list (list): List of molecules (SMILES or RDKit Mol).
+        n_jobs (int): Number of parallel jobs.
+    
+    Returns:
+        Counter: A Counter dictionary mapping fragment SMILES to counts.
+    """
+    fragments = Counter()
+    for mol_frag in mapper(n_jobs)(fragmenter, mol_list):
+        fragments.update(mol_frag)
+    return fragments
+
+
+def compute_scaffolds(mol_list, n_jobs=1, min_rings=2):
+    """
+    Extracts scaffolds from a list of molecules as canonical SMILES.
+    
+    Only scaffolds with at least min_rings rings are considered.
+    
+    Args:
+        mol_list (list): List of molecules.
+        n_jobs (int): Number of parallel jobs.
+        min_rings (int): Minimum number of rings required in a scaffold.
+    
+    Returns:
+        Counter: A Counter mapping scaffold SMILES to counts.
+    """
+    scaffolds = Counter()
+    map_ = mapper(n_jobs)
+    scaffolds = Counter(map_(partial(compute_scaffold, min_rings=min_rings), mol_list))
+    if None in scaffolds:
+        scaffolds.pop(None)
+    return scaffolds
+
+
+def get_n_rings(mol):
+    """
+    Computes the number of rings in a molecule.
+    
+    Args:
+        mol (RDKit Mol): Molecule object.
+    
+    Returns:
+        int: Number of rings.
+    """
+    return mol.GetRingInfo().NumRings()
+
+
+def compute_scaffold(mol, min_rings=2):
+    """
+    Computes the Murcko scaffold of a molecule and returns its canonical SMILES if it has enough rings.
+    
+    Args:
+        mol (str or RDKit Mol): Input molecule.
+        min_rings (int): Minimum number of rings required.
+    
+    Returns:
+        str or None: Canonical SMILES of the scaffold if valid, else None.
+    """
+    mol = get_mol(mol)
+    try:
+        scaffold = MurckoScaffold.GetScaffoldForMol(mol)
+    except (ValueError, RuntimeError):
+        return None
+    n_rings = get_n_rings(scaffold)
+    scaffold_smiles = Chem.MolToSmiles(scaffold)
+    if scaffold_smiles == '' or n_rings < min_rings:
+        return None
+    return scaffold_smiles
+
+
+class Metric:
+    """
+    Abstract base class for chemical metrics.
+    
+    Derived classes should implement the precalc and metric methods.
+    """
+    def __init__(self, n_jobs=1, device='cpu', batch_size=512, **kwargs):
+        self.n_jobs = n_jobs
+        self.device = device
+        self.batch_size = batch_size
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def __call__(self, ref=None, gen=None, pref=None, pgen=None):
+        """
+        Computes the metric between reference and generated molecules.
+        
+        Exactly one of ref or pref, and gen or pgen should be provided.
+        
+        Args:
+            ref: Reference molecule list.
+            gen: Generated molecule list.
+            pref: Precalculated reference metric.
+            pgen: Precalculated generated metric.
+        
+        Returns:
+            Metric value computed by the metric method.
+        """
+        assert (ref is None) != (pref is None), "specify ref xor pref"
+        assert (gen is None) != (pgen is None), "specify gen xor pgen"
+        if pref is None:
+            pref = self.precalc(ref)
+        if pgen is None:
+            pgen = self.precalc(gen)
+        return self.metric(pref, pgen)
+
+    def precalc(self, molecules):
+        """
+        Pre-calculates necessary representations from a list of molecules.
+        Should be implemented by derived classes.
+        """
+        raise NotImplementedError
+
+    def metric(self, pref, pgen):
+        """
+        Computes the metric given precalculated representations.
+        Should be implemented by derived classes.
+        """
+        raise NotImplementedError
+
+
+class FragMetric(Metric):
+    """
+    Metrics based on molecular fragments.
+    """
+    def precalc(self, mols):
+        return {'frag': compute_fragments(mols, n_jobs=self.n_jobs)}
+
+    def metric(self, pref, pgen):
+        return cos_similarity(pref['frag'], pgen['frag'])
+
+
+class ScafMetric(Metric):
+    """
+    Metrics based on molecular scaffolds.
+    """
+    def precalc(self, mols):
+        return {'scaf': compute_scaffolds(mols, n_jobs=self.n_jobs)}
+
+    def metric(self, pref, pgen):
+        return cos_similarity(pref['scaf'], pgen['scaf'])
+
+
+def cos_similarity(ref_counts, gen_counts):
+    """
+    Computes cosine similarity between two molecular vectors.
+    
+    Args:
+        ref_counts (dict): Reference molecular vectors.
+        gen_counts (dict): Generated molecular vectors.
+    
+    Returns:
+        float: Cosine similarity between the two molecular vectors.
+    """
+    if len(ref_counts) == 0 or len(gen_counts) == 0:
+        return np.nan
+    keys = np.unique(list(ref_counts.keys()) + list(gen_counts.keys()))
+    ref_vec = np.array([ref_counts.get(k, 0) for k in keys])
+    gen_vec = np.array([gen_counts.get(k, 0) for k in keys])
+    return 1 - cos_distance(ref_vec, gen_vec)

train.py

@@ -0,0 +1,462 @@
+import os
+import time
+import random
+import pickle
+import argparse
+import os.path as osp
+
+import torch
+import torch.utils.data
+from torch import nn
+from torch_geometric.loader import DataLoader
+
+import wandb
+from rdkit import RDLogger
+
+torch.set_num_threads(5)
+RDLogger.DisableLog('rdApp.*')
+
+from src.util.utils import *
+from src.model.models import Generator, Discriminator, simple_disc
+from src.data.dataset import DruggenDataset
+from src.data.utils import get_encoders_decoders, load_molecules
+from src.model.loss import discriminator_loss, generator_loss
+
+class Train(object):
+    """Trainer for DrugGEN."""
+
+    def __init__(self, config):
+        if config.set_seed:
+            np.random.seed(config.seed)
+            random.seed(config.seed)
+            torch.manual_seed(config.seed)
+            torch.cuda.manual_seed_all(config.seed)
+
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
+
+            os.environ["PYTHONHASHSEED"] = str(config.seed)
+
+            print(f'Using seed {config.seed}')
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
+
+        # Initialize configurations
+        self.submodel = config.submodel
+
+        # Data loader.
+        self.raw_file = config.raw_file  # SMILES containing text file for dataset. 
+                                         # Write the full path to file.
+        self.drug_raw_file = config.drug_raw_file  # SMILES containing text file for second dataset. 
+                                                   # Write the full path to file.
+        
+        # Automatically infer dataset file names from raw file names
+        raw_file_basename = osp.basename(self.raw_file)
+        drug_raw_file_basename = osp.basename(self.drug_raw_file)
+        
+        # Get the base name without extension and add max_atom to it
+        self.max_atom = config.max_atom  # Model is based on one-shot generation.
+        raw_file_base = os.path.splitext(raw_file_basename)[0]
+        drug_raw_file_base = os.path.splitext(drug_raw_file_basename)[0]
+
+        # Change extension from .smi to .pt and add max_atom to the filename
+        self.dataset_file = f"{raw_file_base}{self.max_atom}.pt"
+        self.drugs_dataset_file = f"{drug_raw_file_base}{self.max_atom}.pt"
+
+        self.mol_data_dir = config.mol_data_dir  # Directory where the dataset files are stored.
+        self.drug_data_dir = config.drug_data_dir  # Directory where the drug dataset files are stored.
+        self.dataset_name = self.dataset_file.split(".")[0]
+        self.drugs_dataset_name = self.drugs_dataset_file.split(".")[0]
+        self.features = config.features  # Small model uses atom types as node features. (Boolean, False uses atom types only.)
+                                         # Additional node features can be added. Please check new_dataloarder.py Line 102.
+        self.batch_size = config.batch_size  # Batch size for training.
+        
+        self.parallel = config.parallel
+
+        # Get atom and bond encoders/decoders
+        atom_encoder, atom_decoder, bond_encoder, bond_decoder = get_encoders_decoders(
+            self.raw_file,
+            self.drug_raw_file,
+            self.max_atom
+        )
+        self.atom_encoder = atom_encoder
+        self.atom_decoder = atom_decoder
+        self.bond_encoder = bond_encoder
+        self.bond_decoder = bond_decoder
+
+        self.dataset = DruggenDataset(self.mol_data_dir,
+                                     self.dataset_file,
+                                     self.raw_file,
+                                     self.max_atom,
+                                     self.features,
+                                     atom_encoder=atom_encoder,
+                                     atom_decoder=atom_decoder,
+                                     bond_encoder=bond_encoder,
+                                     bond_decoder=bond_decoder)
+
+        self.loader = DataLoader(self.dataset,
+                                 shuffle=True,
+                                 batch_size=self.batch_size,
+                                 drop_last=True)  # PyG dataloader for the GAN.
+
+        self.drugs = DruggenDataset(self.drug_data_dir, 
+                                 self.drugs_dataset_file, 
+                                 self.drug_raw_file, 
+                                 self.max_atom, 
+                                 self.features,
+                                 atom_encoder=atom_encoder,
+                                 atom_decoder=atom_decoder,
+                                 bond_encoder=bond_encoder,
+                                 bond_decoder=bond_decoder)
+
+        self.drugs_loader = DataLoader(self.drugs, 
+                                       shuffle=True,
+                                       batch_size=self.batch_size, 
+                                       drop_last=True)  # PyG dataloader for the second GAN.
+
+        self.m_dim = len(self.atom_decoder) if not self.features else int(self.loader.dataset[0].x.shape[1]) # Atom type dimension.
+        self.b_dim = len(self.bond_decoder) # Bond type dimension.
+        self.vertexes = int(self.loader.dataset[0].x.shape[0]) # Number of nodes in the graph.
+
+        # Model configurations.
+        self.act = config.act
+        self.lambda_gp = config.lambda_gp
+        self.dim = config.dim
+        self.depth = config.depth
+        self.heads = config.heads
+        self.mlp_ratio = config.mlp_ratio
+        self.ddepth = config.ddepth
+        self.ddropout = config.ddropout
+
+        # Training configurations.
+        self.epoch = config.epoch
+        self.g_lr = config.g_lr
+        self.d_lr = config.d_lr
+        self.dropout = config.dropout
+        self.beta1 = config.beta1
+        self.beta2 = config.beta2
+
+        # Directories.
+        self.log_dir = config.log_dir
+        self.sample_dir = config.sample_dir
+        self.model_save_dir = config.model_save_dir
+
+        # Step size.
+        self.log_step = config.log_sample_step
+
+        # resume training
+        self.resume = config.resume
+        self.resume_epoch = config.resume_epoch
+        self.resume_iter = config.resume_iter
+        self.resume_directory = config.resume_directory
+
+        # wandb configuration
+        self.use_wandb = config.use_wandb
+        self.online = config.online
+        self.exp_name = config.exp_name
+
+        # Arguments for the model.
+        self.arguments = "{}_{}_glr{}_dlr{}_dim{}_depth{}_heads{}_batch{}_epoch{}_dataset{}_dropout{}".format(self.exp_name, self.submodel, self.g_lr, self.d_lr, self.dim, self.depth, self.heads, self.batch_size, self.epoch, self.dataset_name, self.dropout)
+
+        self.build_model(self.model_save_dir, self.arguments)
+
+
+    def build_model(self, model_save_dir, arguments):
+        """Create generators and discriminators."""
+        
+        ''' Generator is based on Transformer Encoder: 
+            
+            @ g_conv_dim: Dimensions for MLP layers before Transformer Encoder
+            @ vertexes: maximum length of generated molecules (atom length)
+            @ b_dim: number of bond types
+            @ m_dim: number of atom types (or number of features used)
+            @ dropout: dropout possibility
+            @ dim: Hidden dimension of Transformer Encoder
+            @ depth: Transformer layer number
+            @ heads: Number of multihead-attention heads
+            @ mlp_ratio: Read-out layer dimension of Transformer
+            @ drop_rate: depricated  
+            @ tra_conv: Whether module creates output for TransformerConv discriminator
+            '''
+        self.G = Generator(self.act,
+                           self.vertexes,
+                           self.b_dim,
+                           self.m_dim,
+                           self.dropout,
+                           dim=self.dim,
+                           depth=self.depth,
+                           heads=self.heads,
+                           mlp_ratio=self.mlp_ratio)
+
+        ''' Discriminator implementation with Transformer Encoder:
+            
+            @ act: Activation function for MLP
+            @ vertexes: maximum length of generated molecules (molecule length)
+            @ b_dim: number of bond types
+            @ m_dim: number of atom types (or number of features used)
+            @ dropout: dropout possibility
+            @ dim: Hidden dimension of Transformer Encoder
+            @ depth: Transformer layer number
+            @ heads: Number of multihead-attention heads
+            @ mlp_ratio: Read-out layer dimension of Transformer'''
+
+        self.D = Discriminator(self.act,
+                                self.vertexes,
+                                self.b_dim,
+                                self.m_dim,
+                                self.ddropout,
+                                dim=self.dim,
+                                depth=self.ddepth,
+                                heads=self.heads,
+                                mlp_ratio=self.mlp_ratio)
+
+        self.g_optimizer = torch.optim.AdamW(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])
+        self.d_optimizer = torch.optim.AdamW(self.D.parameters(), self.d_lr, [self.beta1, self.beta2])
+
+        network_path = os.path.join(model_save_dir, arguments)
+        self.print_network(self.G, 'G', network_path)
+        self.print_network(self.D, 'D', network_path)
+
+        if self.parallel and torch.cuda.device_count() > 1:
+            print(f"Using {torch.cuda.device_count()} GPUs!")
+            self.G = nn.DataParallel(self.G)
+            self.D = nn.DataParallel(self.D)
+
+        self.G.to(self.device)
+        self.D.to(self.device)
+
+    def print_network(self, model, name, save_dir):
+        """Print out the network information."""
+        num_params = 0
+        for p in model.parameters():
+            num_params += p.numel()
+
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+
+        network_path = os.path.join(save_dir, "{}_modules.txt".format(name))
+        with open(network_path, "w+") as file:
+            for module in model.modules():
+                file.write(f"{module.__class__.__name__}:\n")
+                print(module.__class__.__name__)
+                for n, param in module.named_parameters():
+                    if param is not None:
+                        file.write(f"  - {n}: {param.size()}\n")
+                        print(f"  - {n}: {param.size()}")
+                break
+            file.write(f"Total number of parameters: {num_params}\n")
+            print(f"Total number of parameters: {num_params}\n\n")
+
+    def restore_model(self, epoch, iteration, model_directory):
+        """Restore the trained generator and discriminator."""
+        print('Loading the trained models from epoch / iteration {}-{}...'.format(epoch, iteration))
+
+        G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(epoch, iteration))
+        D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(epoch, iteration))
+        self.G.load_state_dict(torch.load(G_path, map_location=lambda storage, loc: storage))
+        self.D.load_state_dict(torch.load(D_path, map_location=lambda storage, loc: storage))
+
+    def save_model(self, model_directory, idx,i):
+        G_path = os.path.join(model_directory, '{}-{}-G.ckpt'.format(idx+1,i+1))
+        D_path = os.path.join(model_directory, '{}-{}-D.ckpt'.format(idx+1,i+1))
+        torch.save(self.G.state_dict(), G_path)
+        torch.save(self.D.state_dict(), D_path)
+
+    def reset_grad(self):
+        """Reset the gradient buffers."""
+        self.g_optimizer.zero_grad()
+        self.d_optimizer.zero_grad()
+
+    def train(self, config):
+        ''' Training Script starts from here'''
+        if self.use_wandb:
+            mode = 'online' if self.online else 'offline'
+        else:
+            mode = 'disabled'
+        kwargs = {'name': self.exp_name, 'project': 'druggen', 'config': config,
+                'settings': wandb.Settings(_disable_stats=True), 'reinit': True, 'mode': mode, 'save_code': True}
+        wandb.init(**kwargs)
+
+        wandb.save(os.path.join(self.model_save_dir, self.arguments, "G_modules.txt"))
+        wandb.save(os.path.join(self.model_save_dir, self.arguments, "D_modules.txt"))
+
+        self.model_directory = os.path.join(self.model_save_dir, self.arguments)
+        self.sample_directory = os.path.join(self.sample_dir, self.arguments)
+        self.log_path = os.path.join(self.log_dir, "{}.txt".format(self.arguments))
+        if not os.path.exists(self.model_directory):
+            os.makedirs(self.model_directory)
+        if not os.path.exists(self.sample_directory):
+            os.makedirs(self.sample_directory)
+
+        # smiles data for metrics calculation.
+        drug_smiles = [line for line in open(self.drug_raw_file, 'r').read().splitlines()]
+        drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
+        drug_vecs = [AllChem.GetMorganFingerprintAsBitVect(x, 2, nBits=1024) for x in drug_mols if x is not None]
+
+        if self.resume:
+            self.restore_model(self.resume_epoch, self.resume_iter, self.resume_directory)
+
+        # Start training.
+        print('Start training...')
+        self.start_time = time.time()
+        for idx in range(self.epoch):
+            # =================================================================================== #
+            #                             1. Preprocess input data                                #
+            # =================================================================================== #
+            # Load the data
+            dataloader_iterator = iter(self.drugs_loader)
+
+            wandb.log({"epoch": idx})
+
+            for i, data in enumerate(self.loader):
+                try:
+                    drugs = next(dataloader_iterator)
+                except StopIteration:
+                    dataloader_iterator = iter(self.drugs_loader)
+                    drugs = next(dataloader_iterator)
+
+                wandb.log({"iter": i})
+
+                # Preprocess both dataset
+                real_graphs, a_tensor, x_tensor = load_molecules(
+                    data=data,
+                    batch_size=self.batch_size,
+                    device=self.device,
+                    b_dim=self.b_dim,
+                    m_dim=self.m_dim,
+                )
+
+                drug_graphs, drugs_a_tensor, drugs_x_tensor = load_molecules(
+                    data=drugs,
+                    batch_size=self.batch_size,
+                    device=self.device,
+                    b_dim=self.b_dim,
+                    m_dim=self.m_dim,
+                )
+
+                # Training configuration.
+                GEN_node = x_tensor             # Generator input node features (annotation matrix of real molecules)
+                GEN_edge = a_tensor             # Generator input edge features (adjacency matrix of real molecules)
+                if self.submodel == "DrugGEN":
+                    DISC_node = drugs_x_tensor  # Discriminator input node features (annotation matrix of drug molecules)
+                    DISC_edge = drugs_a_tensor  # Discriminator input edge features (adjacency matrix of drug molecules)
+                elif self.submodel == "NoTarget":
+                    DISC_node = x_tensor      # Discriminator input node features (annotation matrix of real molecules)
+                    DISC_edge = a_tensor      # Discriminator input edge features (adjacency matrix of real molecules)
+
+                # =================================================================================== #
+                #                                     2. Train the GAN                                #
+                # =================================================================================== #
+
+                loss = {}
+                self.reset_grad()
+                # Compute discriminator loss.
+                node, edge, d_loss = discriminator_loss(self.G,
+                                            self.D,
+                                            DISC_edge,
+                                            DISC_node,
+                                            GEN_edge,
+                                            GEN_node,
+                                            self.batch_size,
+                                            self.device,
+                                            self.lambda_gp)
+                d_total = d_loss
+                wandb.log({"d_loss": d_total.item()})
+
+                loss["d_total"] = d_total.item()
+                d_total.backward()
+                self.d_optimizer.step()
+
+                self.reset_grad()
+
+                # Compute generator loss.
+                generator_output = generator_loss(self.G,
+                                                    self.D,
+                                                    GEN_edge,
+                                                    GEN_node,
+                                                    self.batch_size)
+                g_loss, node, edge, node_sample, edge_sample = generator_output
+                g_total = g_loss
+                wandb.log({"g_loss": g_total.item()})
+
+                loss["g_total"] = g_total.item()
+                g_total.backward()
+                self.g_optimizer.step()
+
+                # Logging.
+                if (i+1) % self.log_step == 0:
+                    logging(self.log_path, self.start_time, i, idx, loss, self.sample_directory,
+                            drug_smiles,edge_sample, node_sample, self.dataset.matrices2mol,
+                            self.dataset_name, a_tensor, x_tensor, drug_vecs)
+
+                    mol_sample(self.sample_directory, edge_sample.detach(), node_sample.detach(),
+                               idx, i, self.dataset.matrices2mol, self.dataset_name)
+                    print("samples saved at epoch {} and iteration {}".format(idx,i))
+
+                    self.save_model(self.model_directory, idx, i)
+                    print("model saved at epoch {} and iteration {}".format(idx,i))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    # Data configuration.
+    parser.add_argument('--raw_file', type=str, required=True)
+    parser.add_argument('--drug_raw_file', type=str, required=False, help='Required for DrugGEN model, optional for NoTarget')
+    parser.add_argument('--drug_data_dir', type=str, default='data')
+    parser.add_argument('--mol_data_dir', type=str, default='data')
+    parser.add_argument('--features', action='store_true', help='features dimension for nodes')
+
+    # Model configuration.
+    parser.add_argument('--submodel', type=str, default="DrugGEN", help="Chose model subtype: DrugGEN, NoTarget", choices=['DrugGEN', 'NoTarget'])
+    parser.add_argument('--act', type=str, default="relu", help="Activation function for the model.", choices=['relu', 'tanh', 'leaky', 'sigmoid'])
+    parser.add_argument('--max_atom', type=int, default=45, help='Max atom number for molecules must be specified.')
+    parser.add_argument('--dim', type=int, default=128, help='Dimension of the Transformer Encoder model for the GAN.')
+    parser.add_argument('--depth', type=int, default=1, help='Depth of the Transformer model from the GAN.')
+    parser.add_argument('--ddepth', type=int, default=1, help='Depth of the Transformer model from the discriminator.')
+    parser.add_argument('--heads', type=int, default=8, help='Number of heads for the MultiHeadAttention module from the GAN.')
+    parser.add_argument('--mlp_ratio', type=int, default=3, help='MLP ratio for the Transformer.')
+    parser.add_argument('--dropout', type=float, default=0., help='dropout rate')
+    parser.add_argument('--ddropout', type=float, default=0., help='dropout rate for the discriminator')
+    parser.add_argument('--lambda_gp', type=float, default=10, help='Gradient penalty lambda multiplier for the GAN.')
+
+    # Training configuration.
+    parser.add_argument('--batch_size', type=int, default=128, help='Batch size for the training.')
+    parser.add_argument('--epoch', type=int, default=10, help='Epoch number for Training.')
+    parser.add_argument('--g_lr', type=float, default=0.00001, help='learning rate for G')
+    parser.add_argument('--d_lr', type=float, default=0.00001, help='learning rate for D')
+    parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for Adam optimizer')
+    parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer')
+    parser.add_argument('--log_dir', type=str, default='experiments/logs')
+    parser.add_argument('--sample_dir', type=str, default='experiments/samples')
+    parser.add_argument('--model_save_dir', type=str, default='experiments/models')
+    parser.add_argument('--log_sample_step', type=int, default=1000, help='step size for sampling during training')
+
+    # Resume training.
+    parser.add_argument('--resume', type=bool, default=False, help='resume training')
+    parser.add_argument('--resume_epoch', type=int, default=None, help='resume training from this epoch')
+    parser.add_argument('--resume_iter', type=int, default=None, help='resume training from this step')
+    parser.add_argument('--resume_directory', type=str, default=None, help='load pretrained weights from this directory')
+
+    # Seed configuration.
+    parser.add_argument('--set_seed', action='store_true', help='set seed for reproducibility')
+    parser.add_argument('--seed', type=int, default=1, help='seed for reproducibility')
+
+    # wandb configuration.
+    parser.add_argument('--use_wandb', action='store_true', help='use wandb for logging')
+    parser.add_argument('--online', action='store_true', help='use wandb online')
+    parser.add_argument('--exp_name', type=str, default='druggen', help='experiment name')
+    parser.add_argument('--parallel', action='store_true', help='Parallelize training')
+
+    config = parser.parse_args()
+
+    # Check if drug_raw_file is provided when using DrugGEN model
+    if config.submodel == "DrugGEN" and not config.drug_raw_file:
+        parser.error("--drug_raw_file is required when using DrugGEN model")
+
+    # If using NoTarget model and drug_raw_file is not provided, use a dummy file
+    if config.submodel == "NoTarget" and not config.drug_raw_file:
+        config.drug_raw_file = "data/akt_train.smi"  # Use a reference file for NoTarget model (AKT) (not used for training for ease of use and encoder/decoder's)
+
+    trainer = Train(config)
+    trainer.train(config)