Acegen rl

.gitignore

@@ -163,5 +163,12 @@ cython_debug/
 results/
 logs_*
 
+# Configs unless specified
+*.yaml
+
+# Nohup files
+*.out
+*.nohup
+
 *.DS_Store*
 *gpt2_enamine_real.ckpt*

acegen/data/chem_utils.py

@@ -2,6 +2,8 @@
 
 import numpy as np
 
+import torch
+
 from rdkit.Chem import AllChem as Chem, Draw
 
 
@@ -60,3 +62,21 @@ def draw(mol_list, molsPerRow=5, subImgSize=(300, 300)):
     image = Draw.MolsToGridImage(mols, molsPerRow=molsPerRow, subImgSize=subImgSize)
 
     return image
+
+
+def get_fp(mol):
+    """Create a Circular/Path based fingerprint from a SMILES string or RDKitMol."""
+    mol = get_mol(mol)
+    if mol:
+        ecfp = Chem.GetMorganFingerprintAsBitVect(mol, radius=3, nBits=256)
+        rdk = Chem.RDKFingerprint(mol, maxPath=6, fpSize=256, nBitsPerHash=2)
+        fp = torch.cat([torch.tensor(ecfp), torch.tensor(rdk)])
+        return fp
+    else:
+        return torch.zeros((512), dtype=torch.int64)
+
+
+def get_fp_hist(mols):
+    """Compute the histogram of fingerprints from a list of SMILES strings or Mols."""
+    fp_hist = torch.vstack([get_fp(mol) for mol in mols]).sum(0)
+    return fp_hist

acegen/data/utils.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import warnings
+
 import torch
 from tensordict import TensorDict
 
@@ -75,10 +77,18 @@ def collate_smiles_to_tensordict(
     """Function to take a list of encoded sequences and turn them into a tensordict."""
     collated_arr = torch.ones(len(arr), max_length) * -1
     for i, seq in enumerate(arr):
-        collated_arr[i, : seq.size(0)] = seq
+        if seq.size(0) > max_length:
+            warnings.warn(
+                f"Sequence {i} is longer than max_length. Truncating to {max_length}."
+            )
+            collated_arr[i, :max_length] = seq[:max_length]
+        else:
+            collated_arr[i, : seq.size(0)] = seq
     data = smiles_to_tensordict(
         collated_arr, reward=reward, replace_mask_value=0, device=device
     )
-    data.set("sequence", data.get("observation"))
-    data.set("sequence_mask", data.get("mask"))
+    data.set("sequence", data.get("observation").clone())
+    data.set("sequence_mask", data.get("mask").clone())
+    data.set(("next", "sequence"), data.get("next", "observation").clone())
+    data.set(("next", "sequence_mask"), data.get("next", "mask").clone())
     return data

acegen/models/__init__.py

@@ -1,5 +1,6 @@
 import logging
 import tarfile
+from functools import partial
 from importlib import import_module, resources
 from pathlib import Path
 
@@ -67,6 +68,14 @@ def extract(path):
         resources.files("acegen.priors") / "gru_chembl_filtered.ckpt",
         SMILESTokenizerChEMBL(),
     ),
+    "gru_chembl34": (
+        create_gru_actor,
+        create_gru_critic,
+        create_gru_actor_critic,
+        resources.files("acegen.priors") / "gru_chembl34_vocabulary.ckpt",
+        resources.files("acegen.priors") / "gru_chembl34.ckpt",
+        SMILESTokenizerChEMBL(),
+    ),
     "lstm": (
         create_lstm_actor,
         create_lstm_critic,
@@ -75,6 +84,17 @@ def extract(path):
         resources.files("acegen.priors") / "lstm_chembl.ckpt",
         SMILESTokenizerChEMBL(),
     ),
+    "lstm_guacamol": (
+        partial(create_lstm_actor, embedding_size=1024, hidden_size=1024, dropout=0.2),
+        partial(create_lstm_critic, embedding_size=1024, hidden_size=1024, dropout=0.2),
+        partial(
+            create_lstm_actor_critic, embedding_size=1024, hidden_size=1024, dropout=0.2
+        ),
+        resources.files("acegen.priors") / "lstm_guacamol_vocabulary.txt",
+        resources.files("acegen.priors")
+        / "lstm_guacamol_model_final_0.473_acegen.ckpt",
+        SMILESTokenizerGuacaMol(),
+    ),
     "gpt2": (
         create_gpt2_actor,
         create_gpt2_critic,

acegen/models/utils.py

@@ -1,4 +1,8 @@
 import warnings
+from typing import Union
+
+import torch
+from tensordict.nn import TensorDictModule
 
 
 def adapt_state_dict(source_state_dict: dict, target_state_dict: dict):
@@ -32,3 +36,15 @@ def adapt_state_dict(source_state_dict: dict, target_state_dict: dict):
         target_state_dict[key_target] = value_source
 
     return target_state_dict
+
+
+def reinitialize_model(
+    model: Union[torch.nn.Module, TensorDictModule], seed: int = 101
+):
+    """Random initialization of a models parameters."""
+    torch.manual_seed(seed)
+    for p in model.parameters():
+        if len(p.shape) == 1:
+            torch.nn.init.constant_(p, 0)
+        else:
+            torch.nn.init.uniform_(p)

acegen/rl_env/baselines.py

@@ -0,0 +1,38 @@
+import torch
+
+
+class MovingAverageBaseline:
+    """Class to keep track on the running mean and variance of tensors batches."""
+
+    def __init__(self, epsilon=1e-3, shape=(), device=torch.device("cpu")):
+        self.mean = torch.zeros(shape, dtype=torch.float64).to(device)
+        self.std = torch.zeros(shape, dtype=torch.float64).to(device)
+        self.count = epsilon
+
+    def update(self, x):
+        batch_mean = torch.mean(x, dim=0)
+        batch_std = torch.std(x, dim=0)
+        batch_count = x.shape[0]
+        self.update_from_moments(batch_mean, batch_std, batch_count)
+
+    def update_from_moments(self, batch_mean, batch_std, batch_count):
+        delta = batch_mean - self.mean
+        std_delta = batch_std - self.std
+        tot_count = self.count + batch_count
+        new_mean = self.mean + delta * batch_count / tot_count
+        new_std = self.std + std_delta * batch_count / tot_count
+        new_count = tot_count
+        self.mean, self.std, self.count = new_mean, new_std, new_count
+
+
+class LeaveOneOutBaseline:
+    """Class to compute the leave-one-out baseline for a given tensor."""
+
+    def __init__(self):
+        self.mean = None
+
+    def update(self, x):
+        with torch.no_grad():
+            loo = x.unsqueeze(1).expand(-1, x.size(0))
+            loo_mask = 1 - torch.eye(loo.size(0), device=loo.device)
+            self.mean = (loo * loo_mask).sum(0) / loo_mask.sum(0)

acegen/rl_env/token_env.py

@@ -3,10 +3,10 @@
 import torch
 from tensordict.tensordict import TensorDict, TensorDictBase
 from torchrl.data import (
-    CompositeSpec,
-    DiscreteTensorSpec,
+    Composite,
+    Categorical,
     OneHotDiscreteTensorSpec,
-    UnboundedContinuousTensorSpec,
+    Unbounded,
 )
 from torchrl.data.utils import DEVICE_TYPING
 from torchrl.envs import EnvBase
@@ -183,9 +183,9 @@ def _set_specs(self) -> None:
         obs_spec = (
             OneHotDiscreteTensorSpec
             if self.one_hot_obs_encoding
-            else DiscreteTensorSpec
+            else Categorical
         )
-        self.observation_spec = CompositeSpec(
+        self.observation_spec = Composite(
             {
                 "observation": obs_spec(
                     n=self.length_vocabulary,
@@ -222,9 +222,9 @@ def _set_specs(self) -> None:
         action_spec = (
             OneHotDiscreteTensorSpec
             if self.one_hot_action_encoding
-            else DiscreteTensorSpec
+            else Categorical
         )
-        self.action_spec = CompositeSpec(
+        self.action_spec = Composite(
             {
                 "action": action_spec(
                     n=self.length_vocabulary,
@@ -233,9 +233,9 @@ def _set_specs(self) -> None:
                 )
             }
         ).expand(self.num_envs)
-        self.reward_spec = CompositeSpec(
+        self.reward_spec = Composite(
             {
-                "reward": UnboundedContinuousTensorSpec(
+                "reward": Unbounded(
                     shape=(1,),
                     dtype=torch.float32,
                     device=self.device,
@@ -244,15 +244,15 @@ def _set_specs(self) -> None:
         ).expand(self.num_envs)
 
         self.done_spec = (
-            CompositeSpec(
+            Composite(
                 {
-                    "done": DiscreteTensorSpec(
+                    "done": Categorical(
                         n=2, dtype=torch.bool, device=self.device
                     ),
-                    "truncated": DiscreteTensorSpec(
+                    "truncated": Categorical(
                         n=2, dtype=torch.bool, device=self.device
                     ),
-                    "terminated": DiscreteTensorSpec(
+                    "terminated": Categorical(
                         n=2, dtype=torch.bool, device=self.device
                     ),
                 }

acegen/vocabulary/tokenizers.py

@@ -97,24 +97,22 @@ def __init__(self, start_token="GO", end_token="EOS"):
         self.REGEXP_ORDER = ["brackets", "brcl"]
         self.start_token = start_token
         self.end_token = end_token
+        self.encode_dict = {
+            "Br": "Y",
+            "Cl": "X",
+            "Si": "A",
+            "Se": "Z",
+            "@@": "R",
+            "se": "E",
+        }
+        self.decode_dict = {v: k for k, v in self.encode_dict.items()}
 
     def tokenize(self, data, with_begin_and_end=False):
         """Tokenizes a SMILES string."""
+        for symbol, token in self.encode_dict.items():
+            data = data.replace(symbol, token)
 
-        def split_by(data, regexps):
-            if not regexps:
-                return list(data)
-            regexp = self.REGEXPS[regexps[0]]
-            splitted = regexp.split(data)
-            tokens = []
-            for i, split in enumerate(splitted):
-                if i % 2 == 0:
-                    tokens += split_by(split, regexps[1:])
-                else:
-                    tokens.append(split)
-            return tokens
-
-        tokens = split_by(data, self.REGEXP_ORDER)
+        tokens = list(data)
         if with_begin_and_end:
             tokens = [self.start_token] + tokens + [self.end_token]
         return tokens
@@ -127,6 +125,9 @@ def untokenize(self, tokens, **kwargs):
                 break
             if token != self.start_token:
                 smi += token
+
+        for symbol, token in self.decode_dict.items():
+            smi = smi.replace(symbol, token)
         return smi