ShayneWierbowski
commented
3 years ago I'm just playing around with this repository myself and was also interested in this question (so I'm not extensively familiar with the code). It'd certainly be better to have an official response here, but from what I can tell you should just need to modify the reconstruct method inside the HierVAE model.
def reconstruct(self, batch):
# I think these should match the output from the preprocess.py script
graphs, tensors, _ = batch
# Reformat as tensors (from numpy arrays?)
tree_tensors, graph_tensors = tensors = make_cuda(tensors)
# Encode batch of compounds
root_vecs, tree_vecs, _, graph_vecs = self.encoder(tree_tensors, graph_tensors)
# Modify the root_recs embeddings? (Not actually sure what's happening between these two steps)
# But the output root_vecs here should be the latent embeddings as far as I can tell (so if you just return
# these instead of running the next step, this should be what you want)
root_vecs, root_kl = self.rsample(root_vecs, self.R_mean, self.R_var, perturb=False)
# Convert the laten embedding(s) back into SMILES string(s)
return self.decoder.decode((root_vecs, root_vecs, root_vecs), greedy=True, max_decode_step=150)If you aren't sure how to get everything in the right format for the input to the reconstruct method, I wrote up a test case to read in a model from the checkpoint, take a given smiles string, encode it, and then decode it.
Again, I'm not 100% confident this is accurate, my attempted encode-decode process did not reproduce the same compound, so it's possible I'm misunderstanding something here.
# Imports / Arg Parser / Functions
# Copied from generate.py preprocess.py and / or hgraph/hgnn.py
from multiprocessing import Pool
import math, random, sys
import pickle
import argparse
from functools import partial
import torch
import numpy
from hgraph import *
import rdkit
def make_cuda(tensors):
tree_tensors, graph_tensors = tensors
make_tensor = lambda x: x if type(x) is torch.Tensor else torch.tensor(x)
tree_tensors = [make_tensor(x).long() for x in tree_tensors[:-1]] + [tree_tensors[-1]]
graph_tensors = [make_tensor(x).long() for x in graph_tensors[:-1]] + [graph_tensors[-1]]
return tree_tensors, graph_tensors
def to_numpy(tensors):
convert = lambda x : x.numpy() if type(x) is torch.Tensor else x
a,b,c = tensors
b = [convert(x) for x in b[0]], [convert(x) for x in b[1]]
return a, b, c
def tensorize(mol_batch, vocab):
x = MolGraph.tensorize(mol_batch, vocab, common_atom_vocab)
return to_numpy(x)
parser = argparse.ArgumentParser()
parser.add_argument('--vocab', required=True)
parser.add_argument('--atom_vocab', default=common_atom_vocab)
parser.add_argument('--model', required=True)
parser.add_argument('--seed', type=int, default=7)
parser.add_argument('--nsample', type=int, default=10000)
parser.add_argument('--rnn_type', type=str, default='LSTM')
parser.add_argument('--hidden_size', type=int, default=250)
parser.add_argument('--embed_size', type=int, default=250)
parser.add_argument('--batch_size', type=int, default=50)
parser.add_argument('--latent_size', type=int, default=32)
parser.add_argument('--depthT', type=int, default=15)
parser.add_argument('--depthG', type=int, default=15)
parser.add_argument('--diterT', type=int, default=1)
parser.add_argument('--diterG', type=int, default=3)
parser.add_argument('--dropout', type=float, default=0.0)
args = parser.parse_args()
# Parse Vocabuary File
vocab = [x.strip("\r\n ").split() for x in open(args.vocab)]
args.vocab = PairVocab(vocab)
# Test Compound to reconstruct
smiles = ['C=Cc1cccc(C(=O)N2CC(c3ccc(F)cc3)C(C)(C)C2)c1']
print("\nINPUT SMILES: {0}\n".format(" ".join(smiles)))
# Convert SMILES String into MolGraph Tree / Graph Tensors
# (See preprocess.py)
o = tensorize(smiles, args.vocab)
batches, tensors, all_orders = o
# Extract pieces we need
tree_tensors, graph_tensors = make_cuda(tensors)
# Load Checkpoint model
model = HierVAE(args)
model.load_state_dict(torch.load(args.model, map_location=torch.device('cpu'))[0])
model.eval()
# Encode compound
root_vecs, tree_vecs, _, graph_vecs = model.encoder(tree_tensors, graph_tensors)
print("\nLATENT_EMBEDDING\n")
print(root_vecs)
print("\n")
# Unsure what this second step does / what the difference between
# the first and second root_vecs values are?
root_vecs, root_kl = model.rsample(root_vecs, model.R_mean, model.R_var, perturb=False)
print("\nLATENT_EMBEDDING_2\n")
print(root_vecs)
print("\n")
# Decode compound
decoded_smiles = model.decoder.decode((root_vecs, root_vecs, root_vecs), greedy=True, max_decode_step=150)
# The decoded and original smiles / compound do not match
# Not sure if this is because something is done wrong or just
# because this compound is one that couldn't be reconstructed
# accurately
print("DECODED SMILES: {0}".format("".join(decoded_smiles)))
michal-pikusa
Hi,
I've successfully trained a generation model on my set of molecules, and I'm able to sample from it with generate.py. However, I was wondering if it's possible to easily extract the latent vectors for a set of input molecules I used as the training set?
I've previously used your JTNN model, and it was easily done with encode_latent function you had in your model class. However, it seems like HierVAE in hgnn.py does not have such a function. Could you point me to any helper functions I would need to invoke to encode a molecule after training and get its latent vector?
Thank you!