<jupyter_start><jupyter_text>Reward FunctionThis section is for defining the reward function.<jupyter_code>from moses.metrics import mol_passes_filters, QED, SA, logP
from moses.metrics.utils import get_n_rings, get_mol


def get_num_rings_6(mol):
    r = mol.GetRingInfo()
    return len([x for x in r.AtomRings() if len(x) > 6])


def penalized_logP(mol_or_smiles, masked=False, default=-5):
    mol = get_mol(mol_or_smiles)
    if mol is None:
        return default
    reward = logP(mol) - SA(mol) - get_num_rings_6(mol)
    if masked and not mol_passes_filters(mol):
        return default
    return reward
! wget https://media.githubusercontent.com/media/molecularsets/moses/master/data/dataset_v1.csv
df = pd.read_csv('dataset_v1.csv')
df = df[df['SPLIT'] == 'train']
df['plogP'] = df['SMILES'].apply(penalized_logP)
df.to_csv('train_plogp_plogpm.csv', index=None)
md = gentrl.MolecularDataset(sources=[
    {'path':'train_plogp_plogpm.csv',
     'smiles': 'SMILES',
     'prob': 1,
     'plogP' : 'plogP',
    }], 
    props=['plogP'])

from torch.utils.data import DataLoader
train_loader = DataLoader(md, batch_size=50, shuffle=True, num_workers=1, drop_last=True)<jupyter_output><empty_output><jupyter_text>Defining the Encoder, Decoder and Descriptors<jupyter_code>enc = gentrl.RNNEncoder(latent_size=50)
dec = gentrl.DilConvDecoder(latent_input_size=50)
latent_descr = 50 * [('c', 20)]
feature_descr = [('c', 20)]<jupyter_output><empty_output><jupyter_text>Training the VAE (pretrain Step)<jupyter_code>model = gentrl.gentrlVAE(enc,
                         dec,
                         train_loader,
                         latent_descr,
                         feature_descr)<jupyter_output><empty_output><jupyter_text>Pytorch-Lightning trainer (For more details refer [this](https://pytorch-lightning.readthedocs.io/en/latest/trainer.html) )- max_epochs: maximum Epochs- gpus: No of gpus (-1 means all of them and 0 means cpu)<jupyter_code>trainer = Trainer(max_epochs=20, gpus=-1)
trainer.fit(model)  # Training the model<jupyter_output>| Name                 | Type           | Params
----------------------------------------------------
0  | enc                  | RNNEncoder     | 888 K 
1  | enc.embs             | Embedding      | 7 K   
2  | enc.rnn              | GRU            | 789 K 
3  | enc.final_mlp        | Sequential     | 91 K  
4  | enc.final_mlp.0      | Linear         | 65 K  
5  | enc.final_mlp.1      | LeakyReLU      | 0     
6  | enc.final_mlp.2      | Linear         | 25 K  
7  | dec                  | DilConvDecoder | 596 K 
8  | dec.latent_fc        | Linear         | 6 K   
9  | dec.input_embeddings | Embedding      | 3 K   
10 | dec.logits_1x1_layer | Conv1d         | 3 K   
11 | dec.parameters       | ParameterList  | 596 K 
12 | lp                   | LP             | 1 M   
13 | lp.parameters        | ParameterList  | 1 M<jupyter_text>Saving the model<jupyter_code>gentrl.save(model, './model')  # Returns true on successful completion<jupyter_output><empty_output><jupyter_text>Reinforcement Learning (train_rl Step)<jupyter_code>model = gentrl.gentrlRL(penalized_logP, 
                        enc,
                        dec,
                        latent_descr,
                        feature_descr,
                       load_model='./model/')
trainer = Trainer(max_epochs=10, gpus=-1)
trainer.fit(model)
gentrl.save(model, './model')<jupyter_output><empty_output><jupyter_text>Sampling Molecules (Sampling step)<jupyter_code>model(10)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12408140192198983, 'fraction_numerical': 0.022328999434708875, 'mean_word_length': 2.539, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 29, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '502849', 'license': 'MIT', 'n_tokens_mistral': 1309, 'n_tokens_neox': 1223, 'n_words': 239, 'path': 'Example.ipynb', 'repo_name': 'icyray/proGENTRL', 'size': '3538'}