add smiles_ahc

main/smiles_ahc/hparams_default.yaml

@@ -0,0 +1,5 @@
+batch_size: 256
+sigma: 120
+topk: 0.25
+learning_rate: 0.0005
+patience: 5

main/smiles_ahc/hparams_tune.yaml

@@ -0,0 +1,16 @@
+name: smiles_ahc
+method: grid
+metric:
+  goal: maximize
+  name: avg_auc
+parameters:
+  batch_size:
+    values: [64, 128, 256]
+  sigma: 
+    values: [30, 60, 90, 120, 250, 500]
+  topk:
+    values: [0.25, 0.5, 0.75]
+  learning_rate:
+    value: 0.0005
+  patience:
+    value: 5

main/smiles_ahc/model/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Morgan Thomas
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

main/smiles_ahc/model/RL.py

@@ -0,0 +1,130 @@
+import os
+from tqdm.auto import tqdm
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+import torch
+from model.model import Model
+from model import utils
+
+
+class ReinforcementLearning:
+    def __init__(self,
+                 device,
+                 agent,
+                 scoring_function,
+                 save_dir,
+                 optimizer,
+                 learning_rate,
+                 is_molscore=True,
+                 freeze=None):
+        # Device
+        self.device = device
+        # Load agent
+        self.agent = Model.load_from_file(file_path=agent, sampling_mode=False, device=device)
+        # Scoring function
+        self.scoring_function = scoring_function
+        self.molscore = is_molscore
+        self.save_dir = save_dir
+        # Optimizer
+        self.optimizer = optimizer(self.agent.network.parameters(), lr=learning_rate)
+        if freeze is not None:
+            self._freeze_network(freeze)
+        self.record = None
+
+    def train(self, n_steps, save_freq):
+        for step in tqdm(range(n_steps), total=n_steps):
+            self._train_step(step=step)
+            # Save the agent weights every few iterations
+            if step % save_freq == 0 and step != 0:
+                self.agent.save(os.path.join(self.save_dir, f'Agent_{step}.ckpt'))
+        # If the entire training finishes, clean up
+        self.agent.save(os.path.join(self.save_dir, f'Agent_{n_steps}.ckpt'))
+        return self.record
+
+    def _freeze_network(self, freeze):
+        n_freeze = freeze * 4 + 1
+        for i, param in enumerate(self.agent.network.parameters()):
+            if i < n_freeze:  # Freeze parameter
+                param.requires_grad = False
+
+    def _update(self, loss, verbose=False):
+        loss = loss.mean()
+        if verbose:
+            print(f'    Loss: {loss.data:.03f}')
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+
+    def _train_step(self, step):
+        raise NotImplementedError
+
+    def _sample_batch(self, batch_size):
+        seqs, smiles, agent_likelihood, probs, log_probs, critic_values = self.agent.sample_sequences_and_smiles(
+            batch_size)
+        return seqs, smiles, agent_likelihood, probs, log_probs, critic_values
+
+    def _score(self, smiles, step):
+        try:
+            scores = self.scoring_function(smiles)
+            scores = utils.to_tensor(scores).to(self.device)
+        except (Exception, BaseException, SystemExit, KeyboardInterrupt) as e:
+            if self.molscore:
+                # If anything fails, save smiles, agent, scoring_function etc.
+                utils.save_smiles(smiles,
+                                  os.path.join(self.save_dir,
+                                               f'failed_{self.scoring_function.step}.smi'))
+                self.agent.save(os.path.join(self.save_dir, f'Agent_{step}.ckpt'))
+                self.scoring_function.write_scores()
+                self.scoring_function._write_temp_state(step=self.scoring_function.step)
+                self.scoring_function.kill_monitor()
+                raise e
+            else:
+                utils.save_smiles(smiles,
+                                  os.path.join(self.save_dir,
+                                               f'failed_{step + 1}.smi'))
+                self.agent.save(os.path.join(self.save_dir, f'Agent_{step}.ckpt'))
+                raise e
+        return scores
+
+
+class AugmentedHillClimb(ReinforcementLearning):
+    _short_name = 'AHC'
+    def __init__(self, device, agent, scoring_function, save_dir, optimizer, learning_rate, is_molscore=True, freeze=None,
+                 prior=None, batch_size=64, sigma=60, topk=0.5, **kwargs):
+        super().__init__(device, agent, scoring_function, save_dir, optimizer, learning_rate, is_molscore, freeze=None)
+
+        # Load prior
+        if prior is None:
+            self.prior = Model.load_from_file(file_path=agent, sampling_mode=True, device=device)
+        else:
+            self.prior = Model.load_from_file(file_path=prior, sampling_mode=True, device=device)
+        # Parameters
+        self.batch_size = batch_size
+        self.sigma = sigma
+        self.topk = topk
+        # Record
+        self.record = {'loss': [],
+                       'prior_nll': [],
+                       'agent_nll': []}
+
+    def _train_step(self, step):
+        # Sample
+        seqs, smiles, agent_likelihood, probs, log_probs, critic_values = self._sample_batch(self.batch_size)
+        # Score
+        scores = self._score(smiles, step)
+        # Compute loss
+        agent_likelihood = - agent_likelihood
+        prior_likelihood = - self.prior.likelihood(seqs)
+        augmented_likelihood = prior_likelihood + self.sigma * scores
+        sscore, sscore_idxs = scores.sort(descending=True)
+        loss = torch.pow((augmented_likelihood - agent_likelihood), 2)
+        # Update
+        self.record['loss'] += list(loss.detach().cpu().numpy())
+        self.record['prior_nll'] += list(-prior_likelihood.detach().cpu().numpy())
+        self.record['agent_nll'] += list(-agent_likelihood.detach().cpu().numpy())
+        loss = loss[sscore_idxs.data[:int(self.batch_size * self.topk)]]
+        self._update(loss, verbose=False)
+
+
+

main/smiles_ahc/model/dataset.py

@@ -0,0 +1,51 @@
+"""
+Implementation of a SMILES dataset from https://github.com/MolecularAI/Reinvent
+"""
+
+import torch
+import torch.utils.data as tud
+
+
+class Dataset(tud.Dataset):
+    """Custom PyTorch Dataset that takes a file containing \n separated SMILES"""
+
+    def __init__(self, smiles_list, vocabulary, tokenizer):
+        self._vocabulary = vocabulary
+        self._tokenizer = tokenizer
+        self._smiles_list = list(smiles_list)
+
+    def __getitem__(self, i):
+        smi = self._smiles_list[i]
+        tokens = self._tokenizer.tokenize(smi)
+        encoded = self._vocabulary.encode(tokens)
+        return torch.tensor(encoded, dtype=torch.long)  # pylint: disable=E1102
+
+    def __len__(self):
+        return len(self._smiles_list)
+
+    @staticmethod
+    def collate_fn(encoded_seqs):
+        """Converts a list of encoded sequences into a padded tensor"""
+        max_length = max([seq.size(0) for seq in encoded_seqs])
+        collated_arr = torch.zeros(len(encoded_seqs), max_length, dtype=torch.long)  # padded with zeroes
+        for i, seq in enumerate(encoded_seqs):
+            collated_arr[i, :seq.size(0)] = seq
+        return collated_arr
+
+
+def calculate_nlls_from_model(model, smiles, batch_size=128):
+    """
+    Calculates NLL for a set of SMILES strings.
+    :param model: Model object.
+    :param smiles: List or iterator with all SMILES strings.
+    :return : It returns an iterator with every batch.
+    """
+    dataset = Dataset(smiles, model.vocabulary, model.tokenizer)
+    _dataloader = tud.DataLoader(dataset, batch_size=batch_size, collate_fn=Dataset.collate_fn, shuffle=True, generator=torch.Generator(device='cuda'))
+
+    def _iterator(dataloader):
+        for batch in dataloader:
+            nlls = model.likelihood(batch.long())
+            yield nlls.data.cpu().numpy()
+
+    return _iterator(_dataloader), len(_dataloader)