{"cells":[{"metadata":{},"cell_type":"markdown","source":"# InChI Tokenisation\n\nHere we explore a couple of tokenisation strategies for InChI strings for the BMS Molecular Translation challenge in an attempt to keep the sequences as short as possible while retaining the molecular information.\n\nShortened sequences will not only speed up training, they reduce memory usage (important for transformer models since memory use scales quadratically with sequence length). Additionally, shorter sequences may also improve model accuracy since the dependcies between tokens are, in general, shorter. Shorter sequences, will however often lead to larger tokeniser vocabularies."},{"metadata":{"trusted":true},"cell_type":"code","source":"import re\nimport random\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom pathlib import Path","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"path = Path(\"/kaggle/input/bms-molecular-translation/train_labels.csv\")\ntrain_labels = pd.read_csv(path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"inchis = train_labels[\"InChI\"].tolist()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Firstly, lets take a subset of the full dataset to speed things up a bit\n\nrandom.seed(1)\nnum_samples = 500000\nrand_inchis = random.sample(inchis, k=num_samples)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check lengths using naive (single character) tokenisation\n\nnaive_inchi_lens = [len(inchi) for inchi in rand_inchis]\n\nprint(f\"Max length: {max(naive_inchi_lens)}\")\nprint(f\"Min length: {min(naive_inchi_lens)}\")\nprint(f\"Avg length: {sum(naive_inchi_lens) / len(naive_inchi_lens)}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.hist(naive_inchi_lens, bins=100)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"InChIs can be very long strings since each molecular structure is required to have a unique InChI. As you can see above, a naive approach to tokenising the InChIs will lead to some very long sequences. Below we show how we can use knowledge of the InChI standard to produce shorter sequences.\n\nFirstly, a regex which splits off the start of the InChI string as well as atoms with two characters."},{"metadata":{"trusted":true},"cell_type":"code","source":"regex = \"InChI=1S|/|[0-9]|-|\\+|,|\\(|\\)|[A-Z][a-z]|.\"\nprog = re.compile(regex)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tokens = [prog.findall(inchi) for inchi in rand_inchis]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Have a look at some examples :)\n\nfor i, ts in enumerate(tokens[:10]):\n    print(rand_inchis[i])\n    print(ts)\n    print()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"regex_inchi_lens = [len(ts) for ts in tokens]\n\nprint(f\"Max length: {max(regex_inchi_lens)}\")\nprint(f\"Min length: {min(regex_inchi_lens)}\")\nprint(f\"Avg length: {sum(regex_inchi_lens) / len(regex_inchi_lens)}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.hist(regex_inchi_lens, bins=100)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can see a small improvement in sequence length, but most of this probably comes from shortening the first 8 characters into one token. Perhaps we can merge 2 and 3 digit numbers together into a single token (with the drawback of having a larger vocabulary)..."},{"metadata":{"trusted":true},"cell_type":"code","source":"regex = \"InChI=1S|/|[0-9]{3}|[0-9]{2}|[0-9]|-|\\+|,|\\(|\\)|[A-Z][a-z]|.\"\nprog = re.compile(regex)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tokens = [prog.findall(inchi) for inchi in rand_inchis]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i, ts in enumerate(tokens[:10]):\n    print(rand_inchis[i])\n    print(ts)\n    print()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"regex_inchi_lens = [len(ts) for ts in tokens]\n\nprint(f\"Max length: {max(regex_inchi_lens)}\")\nprint(f\"Min length: {min(regex_inchi_lens)}\")\nprint(f\"Avg length: {sum(regex_inchi_lens) / len(regex_inchi_lens)}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.hist(naive_inchi_lens, bins=100)\nplt.hist(regex_inchi_lens, bins=100)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We can see a much bigger improvement here, crucially the longest sequences is now only 273 tokens (vs the 386 we had before). But perhaps we can do even better using other molecular notations..."},{"metadata":{},"cell_type":"markdown","source":"### Tokenising SMILES\n\nSMILES (**S**implified **M**olecular-**I**nput **L**ine-**E**ntry **S**ystem)[1] is another commonly used molecular representation, which can be easily constructed from a molecule's InChI string. Similarly, it is straightforward to convert a SMILES to InChI. Unlike InChI strings, however, a molecule can have multiple valid SMILES representations. However, we may be able to use SMILES strings to significantly reduce the number of tokens in a molecular representation.\n\n[1] Weininger, David. \"SMILES, a chemical language and information system. 1. Introduction to methodology and encoding rules.\" Journal of chemical information and computer sciences 28.1 (1988): 31-36."},{"metadata":{},"cell_type":"markdown","source":"First, we need to employ the help of RDKit!"},{"metadata":{"trusted":true},"cell_type":"code","source":"!conda install -y rdkit -c rdkit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from rdkit import Chem","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Lets take a few examples first...\n\nsmall_sample_inchis = random.sample(rand_inchis, k=5)\nmols = [Chem.rdinchi.InchiToMol(inchi)[0] for inchi in small_sample_inchis]\nsmiles = [Chem.MolToSmiles(mol) for mol in mols]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for inchi, smi in zip(rand_inchis, smiles):\n    print(inchi)\n    print(smi)\n    print()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def process_inchi(inchi):\n    mol = Chem.rdinchi.InchiToMol(inchi)[0]\n    smi = Chem.MolToSmiles(mol)\n    return smi","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# This will take a few minutes...\n\nsmiles = [process_inchi(inchi) for inchi in rand_inchis]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Make sure everything worked as expected\n\ninvalids = [smi is None or smi == \"\" for smi in smiles]\nprint(f\"Number of invalid mols: {sum(invalids)}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Now we need a new tokeniser for SMILES\n# Here's one possible tokenisation scheme\n\nsmi_regex = \"\\[|\\]|Br|Cl|\\(|\\)|\\.|=|#|-|\\+|\\\\\\\\|\\/|:|~|@|\\?|>|\\*|\\$|[0-9]{2}|[0-9]|.\"\nsmi_prog = re.compile(smi_regex)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"smi_tokens = [smi_prog.findall(smi) for smi in smiles]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for smi, ts in zip(smiles[:10], smi_tokens[:10]):\n    print(smi)\n    print(ts)\n    print()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"smiles_lens = [len(ts) for ts in smi_tokens]\n\nprint(f\"Max length: {max(smiles_lens)}\")\nprint(f\"Min length: {min(smiles_lens)}\")\nprint(f\"Avg length: {sum(smiles_lens) / len(smiles_lens)}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.hist(naive_inchi_lens, bins=100)\nplt.hist(regex_inchi_lens, bins=100)\nplt.hist(smiles_lens, bins=100)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Using SMILES leads to a very significant reduction in sequence length from the original InChI strings and converting between the representations is straightforward. However, SMILES strings come with the disadvantage that each molecule can be represented in many different ways."},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}