def generate(agent_path,
out,
num=10000,
environ_path='output/RF_cls_ecfp6.pkg'):
""" Generating novel molecules with SMILES representation and
storing them into hard drive as a data frame.
Arguments:
agent_path (str): the neural states file paths for the RNN agent (generator).
out (str): file path for the generated molecules (and scores given by environment).
num (int, optional): the total No. of SMILES that need to be generated. (Default: 10000)
environ_path (str): the file path of the predictor for environment construction.
"""
batch_size = 500
df = pd.DataFrame()
voc = util.Voc("data/voc.txt")
agent = model.Generator(voc)
agent.load_state_dict(torch.load(agent_path))
for i in range(num // batch_size + 1):
if i == 0 and num % batch_size == 0: continue
batch = pd.DataFrame()
samples = agent.sample(batch_size if i != 0 else num % batch_size)
smiles, valids = util.check_smiles(samples, agent.voc)
if environ_path is not None:
# calculating the reward of each SMILES based on the environment (predictor).
environ = util.Environment(environ_path)
scores = environ(smiles)
scores[valids == 0] = 0
valids = scores
batch['SCORE'] = valids
batch['CANONICAL_SMILES'] = smiles
df = df.append(batch)
df.to_csv(out, sep='\t', index=None)
def main():
voc = util.Voc(init_from_file="data/voc_b.txt")
netR_path = 'output/rf_dis.pkg'
netG_path = 'output/net_p'
netD_path = 'output/net_d'
agent_path = 'output/net_gan_%d_%d_%dx%d' % (SIGMA * 10, BL * 10,
BATCH_SIZE, MC)
netR = util.Environment(netR_path)
agent = model.Generator(voc)
agent.load_state_dict(T.load(netG_path + '.pkg'))
df = pd.read_table('data/CHEMBL251.txt')
df = df[df['PCHEMBL_VALUE'] >= 6.5]
data = util.MolData(df, voc)
loader = DataLoader(data,
batch_size=BATCH_SIZE,
shuffle=True,
drop_last=True,
collate_fn=data.collate_fn)
netD = model.Discriminator(VOCAB_SIZE, EMBED_DIM, FILTER_SIZE, NUM_FILTER)
if not os.path.exists(netD_path + '.pkg'):
Train_dis_BCE(netD, agent, loader, epochs=100, out=netD_path)
netD.load_state_dict(T.load(netD_path + '.pkg'))
best_score = 0
log = open(agent_path + '.log', 'w')
for epoch in range(1000):
print('\n--------\nEPOCH %d\n--------' % (epoch + 1))
print('\nPolicy Gradient Training Generator : ')
Train_GAN(agent, netD, netR)
print('\nAdversarial Training Discriminator : ')
Train_dis_BCE(netD, agent, loader, epochs=1)
seqs = agent.sample(1000)
ix = util.unique(seqs)
smiles, valids = util.check_smiles(seqs[ix], agent.voc)
scores = netR(smiles)
scores[valids == False] = 0
unique = (scores >= 0.5).sum() / 1000
if best_score < unique:
T.save(agent.state_dict(), agent_path + '.pkg')
best_score = unique
print("Epoch+: %d average: %.4f valid: %.4f unique: %.4f" %
(epoch, scores.mean(), valids.mean(), unique),
file=log)
for i, smile in enumerate(smiles):
print('%f\t%s' % (scores[i], smile), file=log)
for param_group in agent.optim.param_groups:
param_group['lr'] *= (1 - 0.01)
log.close()
def main():
global Epsilon
# Vocabulary containing all of the tokens for SMILES construction
voc = util.Voc("data/voc.txt")
# File path of predictor in the environment
environ_path = 'output/RF_cls_ecfp6.pkg'
# file path of hidden states in RNN for initialization
initial_path = 'output/net_p'
# file path of hidden states of optimal exploitation network
agent_path = 'output/net_e_%.2f_%.1f_%dx%d' % (Epsilon, Baseline,
BATCH_SIZE, MC)
# file path of hidden states of exploration network
explore_path = 'output/net_p'
# Environment (predictor)
environ = util.Environment(environ_path)
# Agent (generator, exploitation network)
agent = model.Generator(voc)
agent.load_state_dict(torch.load(initial_path + '.pkg'))
# exploration network
explore = model.Generator(voc)
explore.load_state_dict(torch.load(explore_path + '.pkg'))
best_score = 0
log = open(agent_path + '.log', 'w')
for epoch in range(1000):
print('\n--------\nEPOCH %d\n--------' % (epoch + 1))
print('\nForward Policy Gradient Training Generator : ')
Policy_gradient(agent, environ, explore=explore)
seqs = agent.sample(1000)
ix = util.unique(seqs)
smiles, valids = util.check_smiles(seqs[ix], agent.voc)
scores = environ(smiles)
scores[valids == False] = 0
unique = (scores >= 0.5).sum() / 1000
# The model with best percentage of unique desired SMILES will be persisted on the hard drive.
if best_score < unique:
torch.save(agent.state_dict(), agent_path + '.pkg')
best_score = unique
print("Epoch+: %d average: %.4f valid: %.4f unique: %.4f" %
(epoch, scores.mean(), valids.mean(), unique),
file=log)
for i, smile in enumerate(smiles):
print('%f\t%s' % (scores[i], smile), file=log)
# Learing rate exponential decay
for param_group in agent.optim.param_groups:
param_group['lr'] *= (1 - 0.01)
log.close()
def main():
# Construction of the vocabulary
voc = util.Voc("data/voc.txt")
netP_path = 'output/net_pr'
netE_path = 'output/net_ex'
# Pre-training the RNN model with ZINC set
prior = model.Generator(voc)
if not os.path.exists(netP_path + '.pkg'):
print('Exploitation network begins to be trained...')
zinc = util.MolData("data/zinc_corpus.txt", voc, token='SENT')
zinc = DataLoader(zinc,
batch_size=BATCH_SIZE,
shuffle=True,
drop_last=True,
collate_fn=zinc.collate_fn)
prior.fit(zinc, out=netP_path)
print('Exploitation network training is finished!')
prior.load_state_dict(T.load(netP_path + '.pkg'))
# Fine-tuning the RNN model with A2AR set as exploration stragety
explore = model.Generator(voc)
df = pd.read_table('data/chembl_corpus.txt').drop_duplicates(
'CANONICAL_SMILES')
valid = df.sample(BATCH_SIZE)
train = df.drop(valid.index)
explore.load_state_dict(T.load(netP_path + '.pkg'))
# Training set and its data loader
train = util.MolData(train, voc, token='SENT')
train = DataLoader(train,
batch_size=BATCH_SIZE,
collate_fn=train.collate_fn)
# Validation set and its data loader
valid = util.MolData(valid, voc, token='SENT')
valid = DataLoader(valid,
batch_size=BATCH_SIZE,
collate_fn=valid.collate_fn)
print('Exploration network begins to be trained...')
explore.fit(train, loader_valid=valid, out=netE_path, epochs=1000)
print('Exploration network training is finished!')