{"cells":[{"metadata":{},"cell_type":"markdown","source":"<h1>Generate InChI sequences with LSTMs</h1>\n\n<p>Generate InChI sequences with an RNN trained on the ChEMBL smiles Dataset, downloaded from <a href='https://ftp.ebi.ac.uk/pub/databases/chembl/ChEMBLdb/latest/'>here</a> (file <code>chembl_28_chemreps.txt.gz</code>).</p>\n<p>Based on <a href='https://blog.bayeslabs.co/2019/07/04/Generating-Molecules-using-Char-RNN-in-Pytorch.html'>https://blog.bayeslabs.co/2019/07/04/Generating-Molecules-using-Char-RNN-in-Pytorch.html</a>. The links in the article to the complete code are broken, so I've also resorted to this <a href='https://github.com/kevaday/pytorch-char-rnn'>code</a> and this <a href='https://www.youtube.com/watch?v=bbvr-2hY4mE'>video</a>.</p>\n\nPlaying a bit with the hyperparameters (batch size of 512, 4 layers of LSTM, hidden size of 512, embedding dim of 16, ...), I've been able to achieve a loss of around 0.15 and an average score (Levenshtein distance between the target and predicted InChI sequences) of around 25."},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Imports\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.optim import Adam\nfrom torch.nn import LogSoftmax, NLLLoss, Softmax, CrossEntropyLoss\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nimport os\nimport random\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import train_test_split\nimport Levenshtein","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get and uncompress the data\n\n!rm chembl_28_*\n!wget --no-check-certificate https://ftp.ebi.ac.uk/pub/databases/chembl/ChEMBLdb/latest/chembl_28_chemreps.txt.gz\n!gunzip chembl_28_chemreps.txt.gz\n!ls","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Some constants\n\nINPUT_DIR = '.' # where the data lives\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nDATA_COLUMN = 'standard_inchi'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nchemreps = pd.read_csv(os.path.join(INPUT_DIR, 'chembl_28_chemreps.txt'), delimiter='\\t')\nchemreps.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# There are some nulls in the standard_inchi column. Let's drop them\nchemreps = chemreps.dropna()\nchemreps.isnull().sum()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"For this task, we're only interested in the column <code>standard_inchi</code>. Let's create dictionaries for all the characters in this column's values, to map each character to an integer value and back."},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\ncharacters = sorted(list(set(''.join(chemreps[DATA_COLUMN]))))\nnum_characters = len(characters) + 3   # pad, start and end of sequence characters\nint2char = dict(enumerate(characters, 3))\nint2char[0] = ''  # Pad character\nint2char[1] = '[' # Start of sequence character\nint2char[2] = ']' # End of sequence character\nchar2int = {ch:ii for ii,ch in int2char.items()}\nPAD = 0\nSOS = char2int['['] # Start of sequence value\nEOS = char2int[']'] # End of sequence value\nprint(f'there are {num_characters} characters:\\n {char2int}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Let's check the distribution of the lengths\n\nlengths = chemreps[DATA_COLUMN].apply(lambda x: len(x))\npercentiles = [50, 75, 90, 95, 96, 97, 98, 99]\nprint(f'minumum length={lengths.min()} maximum length={lengths.max()}')\nprint(np.percentile(lengths, percentiles))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Let's keep only sequences of 300 or less characters\nSEQ_LENGTH = 403  # Added 3 for PAD, SOS and EOS characters\n\n# Subtract 2 because of the SOS and EOS characters\ndata = chemreps.loc[lengths <= SEQ_LENGTH-2, DATA_COLUMN]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h2>Dataset</h2>"},{"metadata":{"trusted":true},"cell_type":"code","source":"class InChIDataset(Dataset):\n    '''\n    Dataset that generates an array of fixed length of tokens (int) from InChI strings,\n    prefixing and postfixing them with the SOS (start of sequence) and EOS (end of sequence)\n    special characters and padding with the special character PAD till a fixed length if\n    the sequence is smaller.\n    '''\n    \n    def __init__(self, data, seq_length, ctoi):\n        '''\n        Initializes the dataset\n        Args:\n            data: an iterable of InChI strings\n            seq_length: the seq_length of the resulting array of tokens\n            ctoi: a map to converts InChI characters to tokens\n        '''\n        \n        super(Dataset).__init__()\n        \n        x = np.full((len(data), seq_length), PAD, dtype=np.int8) \n        y = np.full((len(data), seq_length), PAD, dtype=np.int8) \n        \n        x[:,0] = SOS\n        \n        for i in range(len(data)):\n            last_ch = min(seq_length - 2, len(data[i])) \n            x[i,1:last_ch+1] = [ctoi[ch] for ch in data[i][:last_ch]]\n            x[i,last_ch+1] = EOS\n        \n        self.x = x\n        y[:,:-1] = x[:,1:] # shift 1 place to the left\n        self.y = y\n    \n    def __len__(self):\n        '''\n        Returns:\n            The number of elements in the dataset\n        '''\n        return self.x.shape[0]\n    \n    def __getitem__(self, idx):\n        '''\n        Args:\n            The index of the dataset element to return\n        Returns:\n            The x and y tensors for element i of the dataset\n        '''\n        return torch.from_numpy(self.x[idx]).long(), torch.from_numpy(self.y[idx]).long()\n    \n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h2>Model</h2>"},{"metadata":{"trusted":true},"cell_type":"code","source":"class InChIRNN(nn.Module):\n    \n    def __init__(self, num_characters, embedding_dim, hidden_size, output_size, num_layers=1, dropout=0):\n        \n        super().__init__()\n        \n        self.num_layers = num_layers\n        self.hidden_size = hidden_size\n        \n        self.embed = nn.Embedding(num_characters, embedding_dim)\n        self.rnn = nn.LSTM(embedding_dim, hidden_size, num_layers, batch_first=True, dropout=dropout)\n        self.dropout = nn.Dropout(dropout)\n        self.fc = nn.Linear(hidden_size, output_size)\n    \n    def forward(self, x, hidden):\n    \n        x_e = self.embed(x)\n        out, hidden = self.rnn(x_e, hidden)\n        \n        # Stack up LSTM outputs using view\n        out = out.contiguous().view(-1, self.hidden_size)     \n        out = self.fc(out)\n        \n        return out, hidden\n    \n    def init_hidden(self, batch_size):\n        \n        hidden = torch.zeros(self.num_layers, batch_size, self.hidden_size)\n        cell = torch.zeros(self.num_layers, batch_size, self.hidden_size)\n        return hidden.to(DEVICE), cell.to(DEVICE)\n    \n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h2>Some utility functions</h2>"},{"metadata":{"trusted":true},"cell_type":"code","source":"def seq2array(seq):\n    '''\n    Convert an InChI sequence of characters to an array of integers, including\n    the SOS and EOS special tokens\n    '''\n    seq_length = len(seq) + 2\n    array = np.zeros((1,seq_length), dtype=np.int8)\n    array[0,0] = SOS\n    array[0,1:-1] = [char2int[ch] for ch in seq]\n    array[0,-1] = EOS\n    return array\n\ndef array2seq(array):\n    '''\n    Convert a sequence of integer tokens to an InChI sequence. The sequence of integers\n    is expected to contain the special tokens SOS and EOS\n    '''\n    array = array.squeeze(0)\n    seq = [int2char[i] for i in array]\n    return ''.join(seq[1:-1])\n\ndef score(y, y_hat):\n    '''\n    Calculate the average Levenshtein score for the distance between y and y_hat\n    Args:\n        y: True labels (dimension batch_size * seq_length) of tokens (int)\n        y_hat: Predicted labels (dimension batch_size * seq_length) of tokens (int)\n    Returns:\n        Average levenshtein distance between the sequences in y and y_hat\n    '''\n    \n    batch_size = y.shape[0]\n    \n    y_str = []\n    y_hat_str = []\n    \n    for i in range(batch_size):\n        y_str.append(''.join([int2char[j] for j in y[i]]))\n        y_hat_str.append(''.join([int2char[j] for j in y_hat[i]]))\n        \n    #print(len(y_str), type(y_str[0]), y_str[0])  \n    sum_levenshtein = 0\n    for i in range(batch_size):\n        sum_levenshtein += Levenshtein.distance(y_str[i], y_hat_str[i])\n    \n    return sum_levenshtein / batch_size","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h2>Training</h2>"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Hyperparameters\n\nEPOCHS = 2\nLR = 1e-3\nBATCH_SIZE = 512\nEMBEDDING_DIM = 16\nHIDDEN_SIZE = 256\nCLIPPING = 2\nNUM_LAYERS = 4\nDROPOUT = 0.1\n\nPRINT_EVERY = 100","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Split data into training and evaluation \n\ntrain, val = train_test_split(data, test_size=0.20)\nprint(f'train: {train.shape}  validation: {val.shape}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Create datasets and data loaders\n\ntrain_dataset = InChIDataset(train.values, SEQ_LENGTH, char2int)\nval_dataset = InChIDataset(val.values, SEQ_LENGTH, char2int)\ntrain_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\nval_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Instantiate model, optimizer and loss function\n\nmodel = InChIRNN(num_characters, EMBEDDING_DIM, HIDDEN_SIZE, num_characters, \n                 num_layers=NUM_LAYERS, dropout=DROPOUT)\noptimizer = Adam(model.parameters(), lr=LR)\ncriterion = CrossEntropyLoss()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Train and validation loop\n\nlr = LR\n\nmodel.to(DEVICE)\n\nstep = 0\nfor epoch in range(EPOCHS):\n    print(f'\\n----------------------- epoch {epoch+1}/{EPOCHS} ------------------------')\n    model.train()\n    running_loss = 0\n    for x, y in tqdm(train_dataloader, \n                     desc='Training loop: ',\n                     total=len(train_dataloader.dataset)//BATCH_SIZE):\n        \n        x, y = x.to(DEVICE), y.to(DEVICE)\n    \n        h0,c0 = model.init_hidden(x.shape[0])\n        output, _ = model(x, (h0,c0))\n        \n        # Calculate cross-entropy loss. The average over all the samples in the batch is returned\n        # (reduction parameter, with default value 'mean')\n        loss = criterion(output.view(-1, num_characters), y.view(-1))\n        current_loss = loss.item() \n        running_loss += current_loss \n        \n        # Backpropagation\n        optimizer.zero_grad()\n        loss.backward()\n        torch.nn.utils.clip_grad_norm_(model.parameters(), CLIPPING)\n        optimizer.step()\n              \n        if step % PRINT_EVERY == 0:\n            print(f'step:{step} loss:{current_loss}')\n            \n        step += 1\n        \n        \n    \n    else:  # end of epoch\n        \n        \n        # Print training loss for the epoch\n        print(f'*** training loss = {running_loss/(int(np.ceil(len(train) // BATCH_SIZE)))}')\n        \n        # Validation loop\n        \n        model.eval()\n        validation_loss = 0\n        sum_levenshtein = 0     # Mean levenshtein distance for each batch\n       \n        with torch.no_grad():\n            \n            \n            for x, y in tqdm(val_dataloader, desc='Validation loop: ', \n                             total=len(val_dataloader.dataset)//BATCH_SIZE):\n        \n                x, y = x.to(DEVICE), y.to(DEVICE)\n    \n                h0,c0 = model.init_hidden(x.shape[0])\n                output, _ = model(x, (h0,c0))\n        \n                # Calculate cross-entropy loss. The average over all the samples in the batch is returned\n                loss = criterion(output.view(-1, num_characters), y.view(-1))\n                current_loss = loss.item() \n                validation_loss += current_loss \n                \n                # Calculate score\n                y_hat = torch.argmax(output.view(-1, SEQ_LENGTH, num_characters), dim=-1)\n                sum_levenshtein += score(y.cpu().numpy(), y_hat.cpu().numpy())\n                 \n            \n        num_batches = int(np.ceil(len(val_dataloader.dataset) // BATCH_SIZE))\n        print(f'*** validation loss = {validation_loss/num_batches} ')\n        print(f'*** validation score = {sum_levenshtein/num_batches}')\n        \n        \n        # Change learning rate every N epochs\n        \n        if epoch % 1 == 0:\n            lr = lr / 5\n            print(f'\\nLearning rate changed lr={lr}')\n            for g in optimizer.param_groups:\n                g['lr'] = lr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Save the last model parameters\n\ntorch.save(model.state_dict(), 'inchi.pth')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h2>Generate some InChI strings</h2>"},{"metadata":{"trusted":true},"cell_type":"code","source":"def generate(model, T=1):\n    model.eval()\n    sequence = []\n    i = 0\n    hidden = model.init_hidden(1)\n    char_idx = SOS\n    while (char_idx != EOS) & (i < SEQ_LENGTH):\n        x = np.array([char_idx]).reshape(1,1)\n        x = torch.from_numpy(x).long().to(DEVICE)\n        output, hidden = model(x, hidden)\n        \n        probs = F.softmax(output/T, dim=1).squeeze()\n        char_idx = torch.multinomial(probs, 1).cpu().item()\n        char = int2char[char_idx]\n        if char_idx != EOS:\n            sequence.append(char)\n        i += 1\n        \n    return ''.join(sequence)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Let's generate some sequences\nfor i in range(10):\n    print(generate(model))\n    print()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}