Added changes

smi-ted/finetune/finetune_classification.py

@@ -28,7 +28,7 @@ def main(config):
     elif config.smi_ted_version == 'v2':
         from smi_ted_large.load import load_smi_ted
 
-    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output)
+    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output, eval=False)
     model.net.apply(model._init_weights)
     print(model.net)
 
@@ -46,6 +46,7 @@ def main(config):
         hparams=config,
         target_metric=config.target_metric,
         seed=config.start_seed,
+        smi_ted_version=config.smi_ted_version,
         checkpoints_folder=config.checkpoints_folder,
         device=device,
         save_every_epoch=bool(config.save_every_epoch),

smi-ted/finetune/finetune_classification_multitask.py

@@ -60,7 +60,7 @@ def main(config):
     elif config.smi_ted_version == 'v2':
         from smi_ted_large.load import load_smi_ted
 
-    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=len(targets))
+    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=len(targets), eval=False)
     model.net.apply(model._init_weights)
     print(model.net)
 
@@ -78,6 +78,7 @@ def main(config):
         hparams=config,
         target_metric=config.target_metric,
         seed=config.start_seed,
+        smi_ted_version=config.smi_ted_version,
         checkpoints_folder=config.checkpoints_folder,
         device=device,
         save_every_epoch=bool(config.save_every_epoch),

smi-ted/finetune/finetune_regression.py

@@ -28,7 +28,7 @@ def main(config):
     elif config.smi_ted_version == 'v2':
         from smi_ted_large.load import load_smi_ted
 
-    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output)
+    model = load_smi_ted(folder=config.model_path, ckpt_filename=config.ckpt_filename, n_output=config.n_output, eval=False)
     model.net.apply(model._init_weights)
     print(model.net)
 
@@ -48,6 +48,7 @@ def main(config):
         hparams=config,
         target_metric=config.target_metric,
         seed=config.start_seed,
+        smi_ted_version=config.smi_ted_version,
         checkpoints_folder=config.checkpoints_folder,
         device=device,
         save_every_epoch=bool(config.save_every_epoch),

smi-ted/finetune/smi_ted_large/load.py

@@ -318,7 +318,7 @@ class Net(nn.Module):
 
 class MoLEncoder(nn.Module):
 
-    def __init__(self, config, n_vocab):
+    def __init__(self, config, n_vocab, eval=False):
         super(MoLEncoder, self).__init__()
 
         # embeddings
@@ -337,7 +337,7 @@ class MoLEncoder(nn.Module):
             # unless we do deterministic_eval here, we will have random outputs
             feature_map=partial(GeneralizedRandomFeatures, 
                                 n_dims=config['num_feats'], 
-                                deterministic_eval=False),
+                                deterministic_eval=eval),
             activation='gelu'
         )
         self.blocks = builder.get()
@@ -361,7 +361,7 @@ class MoLDecoder(nn.Module):
 class Smi_ted(nn.Module):
     """materials.smi-ted-Large 738M Parameters"""
 
-    def __init__(self, tokenizer, config=None):
+    def __init__(self, tokenizer, config=None, eval=False):
         super(Smi_ted, self).__init__()
 
         # configuration
@@ -373,11 +373,11 @@ class Smi_ted(nn.Module):
 
         # instantiate modules
         if self.config:
-            self.encoder = MoLEncoder(self.config, self.n_vocab)
+            self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
             self.decoder = MoLDecoder(self.n_vocab, self.config['max_len'], self.config['n_embd'])
             self.net = Net(self.config['n_embd'], n_output=self.config['n_output'], dropout=self.config['dropout'])
     
-    def load_checkpoint(self, ckpt_path, n_output):
+    def load_checkpoint(self, ckpt_path, n_outputm eval=False):
         # load checkpoint file
         checkpoint = torch.load(ckpt_path, map_location=torch.device('cpu'))
 
@@ -388,7 +388,7 @@ class Smi_ted(nn.Module):
         self._set_seed(self.config['seed'])
 
         # instantiate modules
-        self.encoder = MoLEncoder(self.config, self.n_vocab)
+        self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
         self.decoder = MoLDecoder(self.n_vocab, self.max_len, self.n_embd)
         self.net = Net(self.n_embd, n_output=self.config['n_output'] if 'n_output' in self.config else n_output, dropout=self.config['dropout'])
 
@@ -493,11 +493,12 @@ class Smi_ted(nn.Module):
 def load_smi_ted(folder="./smi_ted_large", 
               ckpt_filename="smi-ted-Large_30.pt",
               vocab_filename="bert_vocab_curated.txt",
-              n_output=1
+              n_output=1,
+              eval=False
               ):
     tokenizer = MolTranBertTokenizer(os.path.join(folder, vocab_filename))
     model = Smi_ted(tokenizer)
-    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output)
+    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output, eval=eval)
     print('Vocab size:', len(tokenizer.vocab))
     print(f'[FINETUNE MODE - {str(model)}]')
     return model

smi-ted/finetune/smi_ted_light/load.py

@@ -318,7 +318,7 @@ class Net(nn.Module):
 
 class MoLEncoder(nn.Module):
 
-    def __init__(self, config, n_vocab):
+    def __init__(self, config, n_vocab, eval=False):
         super(MoLEncoder, self).__init__()
 
         # embeddings
@@ -337,7 +337,7 @@ class MoLEncoder(nn.Module):
             # unless we do deterministic_eval here, we will have random outputs
             feature_map=partial(GeneralizedRandomFeatures, 
                                 n_dims=config['num_feats'], 
-                                deterministic_eval=False),
+                                deterministic_eval=eval),
             activation='gelu'
         )
         self.blocks = builder.get()
@@ -361,7 +361,7 @@ class MoLDecoder(nn.Module):
 class Smi_ted(nn.Module):
     """materials.smi-ted-Light 289M Parameters"""
 
-    def __init__(self, tokenizer, config=None):
+    def __init__(self, tokenizer, config=None, eval=False):
         super(Smi_ted, self).__init__()
 
         # configuration
@@ -373,11 +373,11 @@ class Smi_ted(nn.Module):
 
         # instantiate modules
         if self.config:
-            self.encoder = MoLEncoder(self.config, self.n_vocab)
+            self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
             self.decoder = MoLDecoder(self.n_vocab, self.config['max_len'], self.config['n_embd'])
             self.net = Net(self.config['n_embd'], n_output=self.config['n_output'], dropout=self.config['dropout'])
     
-    def load_checkpoint(self, ckpt_path, n_output):
+    def load_checkpoint(self, ckpt_path, n_output, eval=False):
         # load checkpoint file
         checkpoint = torch.load(ckpt_path, map_location=torch.device('cpu'))
 
@@ -388,7 +388,7 @@ class Smi_ted(nn.Module):
         self._set_seed(self.config['seed'])
 
         # instantiate modules
-        self.encoder = MoLEncoder(self.config, self.n_vocab)
+        self.encoder = MoLEncoder(self.config, self.n_vocab, eval=eval)
         self.decoder = MoLDecoder(self.n_vocab, self.max_len, self.n_embd)
         self.net = Net(self.n_embd, n_output=self.config['n_output'] if 'n_output' in self.config else n_output, dropout=self.config['dropout'])
 
@@ -493,11 +493,12 @@ class Smi_ted(nn.Module):
 def load_smi_ted(folder="./smi_ted_light", 
               ckpt_filename="smi-ted-Light_40.pt",
               vocab_filename="bert_vocab_curated.txt",
-              n_output=1
+              n_output=1,
+              eval=False
               ):
     tokenizer = MolTranBertTokenizer(os.path.join(folder, vocab_filename))
     model = Smi_ted(tokenizer)
-    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output)
+    model.load_checkpoint(os.path.join(folder, ckpt_filename), n_output, eval=eval)
     print('Vocab size:', len(tokenizer.vocab))
     print(f'[FINETUNE MODE - {str(model)}]')
     return model

smi-ted/finetune/trainers.py

@@ -14,6 +14,7 @@ import numpy as np
 import random
 import args
 import os
+import shutil
 from tqdm import tqdm
 
 # Machine Learning
@@ -25,7 +26,7 @@ from utils import RMSE, sensitivity, specificity
 class Trainer:
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='rmse', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='rmse', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         # data
         self.df_train = raw_data[0]
         self.df_valid = raw_data[1]
@@ -39,6 +40,7 @@ class Trainer:
         # config
         self.target_metric = target_metric
         self.seed = seed
+        self.smi_ted_version = smi_ted_version
         self.checkpoints_folder = checkpoints_folder
         self.save_every_epoch = save_every_epoch
         self.save_ckpt = save_ckpt
@@ -115,28 +117,52 @@ class Trainer:
                 # update best loss
                 best_vloss = val_loss
 
-    def evaluate(self):
-        print("\n=====Test Evaluation=====")
-        self._load_checkpoint(self.last_filename)
-        self.model.eval()
-        tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader)
-
-        # show metrics
-        for m in tst_metrics.keys():
-            print(f"[TEST] Evaluation {m.upper()}: {round(tst_metrics[m], 4)}")
-
-        # save predictions
-        pd.DataFrame(tst_preds).to_csv(
-            os.path.join(
-                self.checkpoints_folder, 
-                f'{self.dataset_name}_{self.target if isinstance(self.target, str) else self.target[0]}_predict_test_seed{self.seed}.csv'), 
-            index=False
-        )
+    def evaluate(self, verbose=True):
+        if verbose:
+            print("\n=====Test Evaluation=====")
+
+        if self.smi_ted_version == 'v1':
+            import smi_ted_light.load as load
+        elif self.smi_ted_version == 'v2':
+            import smi_ted_large.load as load
+        else:
+            raise Exception('Please, specify the SMI-TED version: `v1` or `v2`.')
+
+        # copy vocabulary to checkpoint folder
+        if not os.path.exists(os.path.join(self.checkpoints_folder, 'bert_vocab_curated.txt')):
+            smi_ted_path = os.path.dirname(load.__file__)
+            shutil.copy(os.path.join(smi_ted_path, 'bert_vocab_curated.txt'), self.checkpoints_folder)
+
+        # load model for inference
+        model_inf = load.load_smi_ted(
+            folder=self.checkpoints_folder,
+            ckpt_filename=self.last_filename,
+            eval=True,
+        ).to(self.device)
+
+        # set model evaluation mode
+        model_inf.eval()
+
+        # evaluate on test set
+        tst_preds, tst_loss, tst_metrics = self._validate_one_epoch(self.test_loader, model_inf)
+
+        if verbose:
+            # show metrics
+            for m in tst_metrics.keys():
+                print(f"[TEST] Evaluation {m.upper()}: {round(tst_metrics[m], 4)}")
+
+            # save predictions
+            pd.DataFrame(tst_preds).to_csv(
+                os.path.join(
+                    self.checkpoints_folder, 
+                    f'{self.dataset_name}_{self.target if isinstance(self.target, str) else self.target[0]}_predict_test_seed{self.seed}.csv'), 
+                index=False
+            )
 
     def _train_one_epoch(self):
         raise NotImplementedError
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         raise NotImplementedError
 
     def _print_configuration(self):
@@ -203,9 +229,9 @@ class Trainer:
 class TrainerRegressor(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='rmse', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='rmse', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         super().__init__(raw_data, dataset_name, target, batch_size, hparams,
-                         target_metric, seed, checkpoints_folder, save_every_epoch, save_ckpt, device) 
+                         target_metric, seed, smi_ted_version, checkpoints_folder, save_every_epoch, save_ckpt, device) 
 
     def _train_one_epoch(self):
         running_loss = 0.0
@@ -239,11 +265,13 @@ class TrainerRegressor(Trainer):
 
         return running_loss / len(self.train_loader)
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         data_targets = []
         data_preds = []
         running_loss = 0.0
 
+        model = self.model if model is None else model
+
         with torch.no_grad():
             for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair
@@ -251,8 +279,8 @@ class TrainerRegressor(Trainer):
                 targets = targets.clone().detach().to(self.device)
 
                 # Make predictions for this batch
-                embeddings = self.model.extract_embeddings(smiles).to(self.device)
-                predictions = self.model.net(embeddings).squeeze()
+                embeddings = model.extract_embeddings(smiles).to(self.device)
+                predictions = model.net(embeddings).squeeze()
 
                 # Compute the loss
                 loss = self.loss_fn(predictions, targets)
@@ -292,9 +320,9 @@ class TrainerRegressor(Trainer):
 class TrainerClassifier(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='roc-auc', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         super().__init__(raw_data, dataset_name, target, batch_size, hparams,
-                         target_metric, seed, checkpoints_folder, save_every_epoch, save_ckpt, device) 
+                         target_metric, seed, smi_ted_version, checkpoints_folder, save_every_epoch, save_ckpt, device) 
 
     def _train_one_epoch(self):
         running_loss = 0.0
@@ -328,11 +356,13 @@ class TrainerClassifier(Trainer):
 
         return running_loss / len(self.train_loader)
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         data_targets = []
         data_preds = []
         running_loss = 0.0
 
+        model = self.model if model is None else model
+
         with torch.no_grad():
             for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair
@@ -340,8 +370,8 @@ class TrainerClassifier(Trainer):
                 targets = targets.clone().detach().to(self.device)
 
                 # Make predictions for this batch
-                embeddings = self.model.extract_embeddings(smiles).to(self.device)
-                predictions = self.model.net(embeddings).squeeze()
+                embeddings = model.extract_embeddings(smiles).to(self.device)
+                predictions = model.net(embeddings).squeeze()
 
                 # Compute the loss
                 loss = self.loss_fn(predictions, targets.long())
@@ -397,9 +427,9 @@ class TrainerClassifier(Trainer):
 class TrainerClassifierMultitask(Trainer):
 
     def __init__(self, raw_data, dataset_name, target, batch_size, hparams,
-                 target_metric='roc-auc', seed=0, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
+                 target_metric='roc-auc', seed=0, smi_ted_version=None, checkpoints_folder='./checkpoints', save_every_epoch=False, save_ckpt=True, device='cpu'):
         super().__init__(raw_data, dataset_name, target, batch_size, hparams,
-                         target_metric, seed, checkpoints_folder, save_every_epoch, save_ckpt, device)
+                         target_metric, seed, smi_ted_version, checkpoints_folder, save_every_epoch, save_ckpt, device)
 
     def _prepare_data(self):
         # normalize dataset
@@ -464,12 +494,14 @@ class TrainerClassifierMultitask(Trainer):
 
         return running_loss / len(self.train_loader)
 
-    def _validate_one_epoch(self, data_loader):
+    def _validate_one_epoch(self, data_loader, model=None):
         data_targets = []
         data_preds = []
         data_masks = []
         running_loss = 0.0
 
+        model = self.model if model is None else model
+
         with torch.no_grad():
             for idx, data in enumerate(pbar := tqdm(data_loader)):
                 # Every data instance is an input + label pair + mask
@@ -477,8 +509,8 @@ class TrainerClassifierMultitask(Trainer):
                 targets = targets.clone().detach().to(self.device)
 
                 # Make predictions for this batch
-                embeddings = self.model.extract_embeddings(smiles).to(self.device)
-                predictions = self.model.net(embeddings, multitask=True).squeeze()
+                embeddings = model.extract_embeddings(smiles).to(self.device)
+                predictions = model.net(embeddings, multitask=True).squeeze()
                 predictions = predictions * target_masks.to(self.device)
 
                 # Compute the loss