<jupyter_start><jupyter_text>JTNN quickstart First, install the package if you haven't already. If you use conda:<jupyter_code>! conda create -n jtnn_env --file conda_list.txt<jupyter_output><empty_output><jupyter_text>Then<jupyter_code>! pip install -e .
import os
import pickle

# Disable CUDA (workaround for GPU memory leak issue)
os.environ["CUDA_VISIBLE_DEVICES"]=""

import tqdm
import pandas as pd

from IPython.display import display
from fast_jtnn.fp_calculator import FingerprintCalculator
from fast_jtnn.mol_tree import main_mol_tree
from fast_molvae.preprocess import create_tensor_pickle
from fast_molvae.vae_train import main_vae_train
from rdkit import Chem
from sklearn.cluster import KMeans
# Verify that rdkit version is 2020.09.3 as
# version 2021.03.1 does not seem to work
import rdkit
rdkit.__version__<jupyter_output><empty_output><jupyter_text>Optional: Remove molecules with high-valence atoms<jupyter_code>dataset = open("data/full_train.txt", "r").read().split()
with open("data/full_train.txt", "w") as f:
    for smiles in tqdm.tqdm(dataset):
        mol = Chem.MolFromSmiles(smiles)
        for atom in mol.GetAtoms():
            if atom.GetDegree() > 6:
                print(f"Rejecting high-valence molecule {smiles}")
                break
        else:
            f.write(smiles + "\n")<jupyter_output><empty_output><jupyter_text>Generate vocabulary<jupyter_code>main_mol_tree('data/full_train.txt', 'data/vocab_full.txt')<jupyter_output><empty_output><jupyter_text>Tensorize training set molecules<jupyter_code>create_tensor_pickle('data/full_train.txt', 'data/tensors/tensors_full.p')<jupyter_output><empty_output><jupyter_text>Train VAE<jupyter_code>model = main_vae_train('data/tensors/', 'data/vocab_full.txt', 'data/models', num_workers=4)<jupyter_output><empty_output><jupyter_text>Check that it works<jupyter_code>smiles_list = open("data/full_train.txt", "r").read().split()[:5000]
fp_calculator = FingerprintCalculator("data/models/model.iter-20000", "data/vocab_full.txt")
smiles_list
%pdb on
fps = fp_calculator(smiles_list)
kmeans = KMeans(n_clusters=500, random_state=0).fit(fps)
labels = kmeans.labels_
db = pd.DataFrame()
db['smiles']=smiles_list
db['label']=labels
db
for smiles in db[db['label']==10]['smiles'].values:
    display(Chem.Draw.MolToImage(Chem.MolFromSmiles(smiles)))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11834319526627218, 'fraction_numerical': 0.013102282333051564, 'mean_word_length': 6.194528875379939, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '145181', 'license': 'MIT', 'n_tokens_mistral': 906, 'n_tokens_neox': 849, 'n_words': 159, 'path': 'jtnn_vae_quickstart.ipynb', 'repo_name': 'croningp/JTNN-VAE', 'size': '2366'}