Upload folder using huggingface_hub

README.md

@@ -0,0 +1,70 @@
+
+# De Novo Drug Generator - RNN-VAE
+
+De Novo Drug Generator - RNN-VAE is a deep learning model designed for generating novel drug molecules.
+Training data from ChemBL library
+
+Full project file at https://github.com/teohyc/drug_agent
+
+
+## Usage 
+
+```python
+from rdkit import Chem
+from rdkit.Chem import Draw, Descriptors
+from tree_rnn_vae_infer import generate_candidate_mol
+from tree_rnn_vae_model import TreeEncoder, LatentHead, TreeVAE, TreeDecoder
+
+
+def compute_molecule_props(mol):
+    return {
+        "MW": Descriptors.MolWt(mol),
+        "logP": Descriptors.MolLogP(mol),
+        "HBD": Descriptors.NumHDonors(mol),
+        "HBA": Descriptors.NumHAcceptors(mol),
+    }
+
+
+# display molecule
+def render_molecule_grid(selected):
+    if not selected:
+        return
+
+    mols, legends = [], []
+
+    if isinstance(selected, dict):
+        iterable = selected.items()
+    else:
+        iterable = enumerate(selected, 1)
+
+    for i, item in iterable:
+        if isinstance(selected, dict):
+            smi, props = i, item
+        else:
+            smi, props = item, None
+
+        mol = Chem.MolFromSmiles(smi)
+        if mol:
+            mols.append(mol)
+            if props is None:
+                props = compute_molecule_props(mol)
+            legends.append(
+                f"M{i}
+MW={props['MW']:.0f}, logP={props['logP']:.2f}, "
+                f"HBD={props['HBD']}, HBA={props['HBA']}"
+            )
+
+    img = Draw.MolsToGridImage(
+        mols,
+        molsPerRow=3,
+        subImgSize=(400, 400),
+        legends=legends,
+        useSVG=False,
+    )
+    return img
+
+
+molecules = generate_candidate_mol(num_samples=6, max_len=20) #change to your desired molecule size and number
+img = render_molecule_grid(molecules)
+img.show()
+```

dict_idx2frag.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b45ac2d4f65aa76bd54a6742c48cc5b749d58fedb2f0309c648a258481155e5d
+size 18769

tree_rnn-vae.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5abe40c97c5915ac21890b44d309edba62c9933fc66d8c9e0043644ef46e7cc8
+size 3963609

tree_rnn_vae_infer.py

@@ -0,0 +1,141 @@
+import pandas as pd
+import torch
+from torch import nn
+from torch.utils.data import Dataset, DataLoader
+import math
+import torch
+import torch.nn.functional as F
+from rdkit import Chem
+from rdkit.Chem import Draw, AllChem
+import random
+import re
+from tree_rnn_vae_model import TreeEncoder, LatentHead, TreeVAE, TreeDecoder
+
+def decode_fragments_to_smiles(batch_indices, vocab):
+    smiles_list = []
+    for seq in batch_indices:
+        fragments = [vocab[idx] for idx in seq.tolist()]
+        
+        try:
+            mol = Chem.MolFromSmiles(fragments)
+            if mol is not None:
+                smiles_list.append(Chem.MolToSmiles(mol))
+            else:
+                smiles_list.append(fragments)  # fallback
+        except:
+            smiles_list.append(fragments)
+    return smiles_list
+
+def normalize_attachment_points(smiles):
+    return re.sub(r'\[\d+\*\]', '[*]', smiles)
+
+def deduplicate_fragments(fragments):
+    seen = set()
+    unique = []
+    
+    for f in fragments:
+        if f not in seen:
+            unique.append(f)
+            seen.add(f)
+    return unique
+
+def fragments_to_mols(fragments):
+    mols = []
+    for f in fragments:
+        mol = Chem.MolFromSmiles(f)
+        if mol is None:
+            continue  # skip invalid fragment
+        mols.append(mol)
+    return mols
+
+def get_dummy_atom_indices(mol):
+    return [atom.GetIdx() for atom in mol.GetAtoms() if atom.GetSymbol() == '*']
+
+def get_linear_connection_pairs(mols):
+    """Connect last dummy atom of mol i to first dummy atom of mol i+1"""
+    pairs = []
+    for i in range(len(mols)-1):
+        idx_last_i = get_dummy_atom_indices(mols[i])[-1]
+        idx_first_next = get_dummy_atom_indices(mols[i+1])[0]
+        pairs.append((i, idx_last_i, i+1, idx_first_next))
+    return pairs
+
+def combine_fragments_and_connect(fragments, connection_pairs):
+    mols = fragments_to_mols(fragments)
+    if not mols:
+        return None
+
+    # Combine molecules sequentially
+    combined = Chem.CombineMols(mols[0], mols[1]) if len(mols) > 1 else mols[0]
+    offsets = [0, mols[0].GetNumAtoms()]
+    for i in range(2, len(mols)):
+        combined = Chem.CombineMols(combined, mols[i])
+        offsets.append(sum([m.GetNumAtoms() for m in mols[:i]]))
+
+    rw_mol = Chem.RWMol(combined)
+
+    # Add bonds for connection pairs
+    for f1, a1, f2, a2 in connection_pairs:
+        rw_mol.AddBond(offsets[f1] + a1, offsets[f2] + a2, Chem.BondType.SINGLE)
+
+    # Sanitize
+    try:
+        Chem.SanitizeMol(rw_mol)
+        return Chem.MolToSmiles(rw_mol)
+    except:
+        return None
+
+
+def process_one_molecule(raw_fragments):
+    # Normalize & deduplicate
+    frags = [normalize_attachment_points(f) for f in raw_fragments]
+    frags = deduplicate_fragments(frags)
+
+    if len(frags) == 0:
+        return None
+
+    # Convert to mols
+    mols = fragments_to_mols(frags)
+    if len(mols) < 1:
+        return None
+
+    # Build connections (linear for now)
+    connection_pairs = get_linear_connection_pairs(mols)
+
+    # Combine & sanitize
+    smiles = combine_fragments_and_connect(frags, connection_pairs)
+    return smiles
+
+def replace_dummy_with_carbon(smiles):
+    return smiles.replace('*', 'C')
+    
+
+#main
+def generate_candidate_mol(num_samples=6, max_len=20):
+    '''This function generates candidate molecules using an in-house-designed Tree-RNN VAE model. It returns a list of SMILES strings representing the generated molecules as well as displaying their molecular structure.'''
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    idx2frag = torch.load("dict_idx2frag.pt")      # {index: fragment_string}
+    vocab_size = len(idx2frag)
+
+    model = torch.load("tree_rnn-vae.pt", weights_only=False, map_location=device)
+
+    model.eval()
+    z = torch.randn(num_samples, model.decoder.z_to_hidden.in_features, device=device)
+    sampled_indices = model.sample_from_z(z, sos_idx=None, max_len=max_len)  
+
+    smiles_out = decode_fragments_to_smiles(sampled_indices, idx2frag)
+
+    smiles_list = []
+    for mol in smiles_out:
+        smiles = process_one_molecule(mol)
+        smiles_list.append(smiles)
+
+    smiles_list = [replace_dummy_with_carbon(s) for s in smiles_list]
+    
+    '''mols = [Chem.MolFromSmiles(s) for s in smiles_list]
+    img = Draw.MolsToGridImage(mols, molsPerRow=2, subImgSize=(600,600), legends=smiles_list)
+    img.show()'''
+
+    return smiles_list
+

tree_rnn_vae_model.py

@@ -0,0 +1,132 @@
+import pandas as pd
+import torch
+from torch import nn
+from torch.utils.data import Dataset, DataLoader
+import math
+import torch
+import torch.nn.functional as F
+from rdkit import Chem
+from rdkit.Chem import Draw, AllChem
+import random
+
+class TreeEncoder(nn.Module):
+    def __init__(self, vocab_size, embed_dim, enc_hidden, pad_idx=0):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
+        self.gru = nn.GRU(embed_dim, enc_hidden, batch_first=True)
+        self.enc_hidden = enc_hidden
+
+    def forward(self, x, lengths):
+        """
+        x: [B, L] LongTensor (padded fragment indices)
+        lengths: [B] LongTensor
+        returns: last_hidden: [B, enc_hidden]
+        """
+        emb = self.embed(x)  # [B, L, E]
+        # pack
+        packed = nn.utils.rnn.pack_padded_sequence(emb, lengths.cpu(), batch_first=True, enforce_sorted=False)
+        packed_out, h_n = self.gru(packed)  # h_n: [1, B, enc_hidden]
+        return h_n.squeeze(0)  # [B, enc_hidden]
+    
+class LatentHead(nn.Module):
+    def __init__(self, enc_hidden, z_dim):
+        super().__init__()
+        self.linear_mu = nn.Linear(enc_hidden, z_dim)
+        self.linear_logvar = nn.Linear(enc_hidden, z_dim)
+
+    def forward(self, h):
+        mu = self.linear_mu(h)
+        logvar = self.linear_logvar(h)
+        return mu, logvar
+    
+def reparameterize(mu, logvar):
+    std = (0.5 * logvar).exp()
+    eps = torch.randn_like(std)
+    return mu + eps * std
+
+class TreeDecoder(nn.Module):
+    def __init__(self, vocab_size, embed_dim, dec_hidden, z_dim, pad_idx=0):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=pad_idx)
+        self.z_to_hidden = nn.Linear(z_dim, dec_hidden)
+        # GRU input: embed_dim + z_dim
+        self.gru = nn.GRU(embed_dim + z_dim, dec_hidden, batch_first=True)
+        self.out = nn.Linear(dec_hidden, vocab_size)
+
+    def init_hidden_from_z(self, z):
+        # z: [B, z_dim] -> [1, B, dec_hidden]
+        return torch.tanh(self.z_to_hidden(z)).unsqueeze(0)
+
+    def forward(self, inputs, z, hidden=None):
+        # Full sequence forward (teacher forcing)
+        B, L = inputs.size()
+        if hidden is None:
+            hidden = self.init_hidden_from_z(z)
+
+        emb = self.embed(inputs)                               # [B,L,E]
+        z_exp = z.unsqueeze(1).expand(-1, L, -1)         # [B,L,z_dim]
+        gru_input = torch.cat([emb, z_exp], dim=-1)      # [B,L,E+z]
+        out, hidden = self.gru(gru_input, hidden)        # [B,L,H]
+        logits = self.out(out)                            # [B,L,vocab]
+        return logits, hidden
+
+    def step(self, input_token, z, hidden=None):
+        """
+        Single time-step for autoregressive inference.
+        input_token: [B] long
+        z: [B, z_dim]
+        hidden: previous hidden
+        returns logits [B, vocab], new_hidden
+        """
+        emb = self.embed(input_token).unsqueeze(1)  # [B,1,E]
+        gru_in = torch.cat([emb, z.unsqueeze(1)], dim=-1)  # [B,1,E+z]
+        out, h_n = self.gru(gru_in, hidden)  # out: [B,1,dec_hidden]
+        logits = self.out(out.squeeze(1))  # [B, vocab]
+        return logits, h_n
+    
+class TreeVAE(nn.Module):
+    def __init__(self, vocab_size, embed_dim, enc_hidden, dec_hidden, z_dim, pad_idx=0):
+        super().__init__()
+        self.encoder = TreeEncoder(vocab_size, embed_dim, enc_hidden, pad_idx)
+        self.latent = LatentHead(enc_hidden, z_dim)
+        self.decoder = TreeDecoder(vocab_size, embed_dim, dec_hidden, z_dim, pad_idx)
+
+    def forward(self, x, lengths, tf_prob=1.0):
+        """
+        x: [B, L] padded target sequences (we will use teacher forcing)
+        lengths: [B] actual lengths
+        tf_prob: teacher forcing probability (0..1)
+        Returns: logits [B, L, V], mu, logvar
+        """
+        h_enc = self.encoder(x, lengths)                    # [B, enc_hidden]
+        mu, logvar = self.latent(h_enc)                     # [B, z_dim]
+        z = reparameterize(mu, logvar).to(device)           # [B, z_dim]
+
+        # Decoder with teacher forcing:
+        # For teacher forcing we input the target sequence as inputs (shifted if you want SOS)
+        _, hidden = self.decoder(x, z)                       # predict tokens given the inputs (simpler)
+        return mu, logvar, z, hidden
+
+    def sample_from_z(self, z, sos_idx=None, max_len=32):
+        """
+        z: [B, z_dim] latent
+        returns: generated indices list (B x <=max_len)
+        """
+        if sos_idx is None:
+            sos_idx = random.choice([17, 9, 5, 11, 2]) #if not specified use top 5 most frequent starting fragments
+        self.eval()
+        B = z.size(0)
+        with torch.no_grad():
+            
+            # start token: we will assume user uses a special sos index; if none, use first fragment in vocab
+            input_tok = torch.full((B,), sos_idx, dtype=torch.long, device=z.device)  # [B]
+            hidden = None
+            generated = [input_tok.unsqueeze(1)]
+            for t in range(max_len):
+                logits, hidden = self.decoder.step(input_tok, z, hidden)  # [B, vocab]
+                probs = F.softmax(logits, dim=-1)
+                # sample or argmax; use sampling to get diverse outputs
+                input_tok = torch.multinomial(probs, num_samples=1).squeeze(1)  # [B]
+                generated.append(input_tok.unsqueeze(1))
+            gen = torch.cat(generated, dim=1)  # [B, max_len]
+        return gen  # indices