{ "cells": [ { "cell_type": "markdown", "id": "fa328603-9e2b-4643-8500-ec11c51b5223", "metadata": {}, "source": [ "## Imports and setup" ] }, { "cell_type": "code", "execution_count": 3, "id": "7716fb32-a805-4888-9dac-da4cff4f6e40", "metadata": {}, "outputs": [], "source": [ "import os\n", "import huggingface_hub" ] }, { "cell_type": "code", "execution_count": 4, "id": "432e1636", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "The token has not been saved to the git credentials helper. Pass `add_to_git_credential=True` in this function directly or `--add-to-git-credential` if using via `huggingface-cli` if you want to set the git credential as well.\n", "Token is valid (permission: write).\n", "Your token has been saved to /share/kuleshov/yzs2/discrete-guidance/.hf_cache/token\n", "Login successful\n" ] } ], "source": [ "if os.path.exists(os.path.join(os.environ['HF_HOME'], 'token')):\n", " with open(os.path.join(os.environ['HF_HOME'], 'token'), 'r') as f:\n", " token = f.read().strip()\n", "else:\n", " token = None\n", "huggingface_hub.login(token=token)" ] }, { "cell_type": "code", "execution_count": 26, "id": "e22e86ae", "metadata": {}, "outputs": [], "source": [ "import json\n", "import re\n", "import typing\n", "\n", "import datasets\n", "import numpy as np\n", "import pandas as pd\n", "import rdkit\n", "import transformers\n", "from rdkit import Chem as rdChem\n", "from tqdm.auto import tqdm" ] }, { "cell_type": "code", "execution_count": 7, "id": "aaa00828", "metadata": {}, "outputs": [], "source": [ "# TODO: Update to 2024.03.6 release when available instead of suppressing warning!\n", "# See: https://github.com/rdkit/rdkit/issues/7625#\n", "rdkit.rdBase.DisableLog('rdApp.warning')" ] }, { "cell_type": "markdown", "id": "0a878a71-d33f-43fe-955d-4250950b1eec", "metadata": { "jp-MarkdownHeadingCollapsed": true }, "source": [ "## Create dataset" ] }, { "cell_type": "code", "execution_count": null, "id": "26856fe2", "metadata": {}, "outputs": [], "source": [ "def count_rings_and_bonds(\n", " mol: rdChem.Mol\n", ") -> typing.Dict[str, int]:\n", " \"\"\"Counts bond and ring (by type).\"\"\"\n", " \n", " # Counting rings\n", " ssr = rdChem.GetSymmSSSR(mol)\n", " ring_count = len(ssr)\n", " \n", " ring_sizes = {}\n", " for ring in ssr:\n", " ring_size = len(ring)\n", " if ring_size not in ring_sizes:\n", " ring_sizes[ring_size] = 0\n", " ring_sizes[ring_size] += 1\n", " \n", " # Counting bond types\n", " bond_counts = {\n", " 'single': 0,\n", " 'double': 0,\n", " 'triple': 0,\n", " 'aromatic': 0\n", " }\n", " \n", " for bond in mol.GetBonds():\n", " if bond.GetIsAromatic():\n", " bond_counts['aromatic'] += 1\n", " elif bond.GetBondType() == rdChem.BondType.SINGLE:\n", " bond_counts['single'] += 1\n", " elif bond.GetBondType() == rdChem.BondType.DOUBLE:\n", " bond_counts['double'] += 1\n", " elif bond.GetBondType() == rdChem.BondType.TRIPLE:\n", " bond_counts['triple'] += 1\n", " result = {\n", " 'ring_count': ring_count,\n", " }\n", " for k, v in ring_sizes.items():\n", " result[f\"R{k}\"] = v\n", "\n", " for k, v in bond_counts.items():\n", " result[f\"{k}_bond\"] = v\n", " return result" ] }, { "cell_type": "code", "execution_count": null, "id": "fbde53f7", "metadata": {}, "outputs": [], "source": [ "\"\"\"\n", " Download data and validation indices from:\n", " \"Score-based Generative Modeling of Graphs via the System of Stochastic Differential Equations\"\n", " https://github.com/harryjo97/GDSS\n", " > wget wget https://raw.githubusercontent.com/harryjo97/GDSS/master/data/zinc250k.csv\n", " > wget https://raw.githubusercontent.com/harryjo97/GDSS/master/data/valid_idx_zinc250k.json\n", "\"\"\"\n", "df = pd.read_csv('/Users/yairschiff/Downloads/zinc250k.csv', index_col=0, encoding='utf_8')\n", "feats = []\n", "for i, row in tqdm(df.iterrows(), total=len(df), desc='RDKit feats', leave=False):\n", " feat = {'smiles': row['smiles']}\n", " feat['canonical_smiles'] = rdChem.CanonSmiles(feat['smiles'])\n", " m = rdChem.MolFromSmiles(feat['canonical_smiles'])\n", " feat.update(count_rings_and_bonds(m))\n", " feats.append(feat)\n", "df = pd.merge(df, pd.DataFrame.from_records(feats), on='smiles')\n", "df = df.fillna(0)\n", "for col in df.columns: # recast ring counts as int\n", " if re.search(\"^R[0-9]+$\", col) is not None:\n", " df[col] = df[col].astype(int)\n", "# Re-order columns\n", "df = df[\n", " ['smiles', 'logP', 'qed', 'SAS', 'canonical_smiles',\n", " 'single_bond', 'double_bond', 'triple_bond', 'aromatic_bond',\n", " 'ring_count','R3', 'R4', 'R5', 'R6', 'R7', 'R8', 'R9', 'R10', 'R12', 'R13', 'R14', 'R15', 'R18', 'R24']]" ] }, { "cell_type": "code", "execution_count": null, "id": "1e2d5955", "metadata": {}, "outputs": [], "source": [ "# Read in validation indices\n", "with open('/Users/yairschiff/Downloads/valid_idx_zinc250k.json', 'r') as f:\n", " valid_idxs = json.load(f)\n", "df['validation'] = df.index.isin(valid_idxs).astype(int)" ] }, { "cell_type": "code", "execution_count": null, "id": "2b89b732", "metadata": {}, "outputs": [], "source": [ "# Create HF dataset\n", "dataset = datasets.DatasetDict({\n", " 'train': datasets.Dataset.from_pandas(df[df['validation'] == 0].drop(columns=['validation'])),\n", " 'validation': datasets.Dataset.from_pandas(df[df['validation'] == 1].drop(columns=['validation'])),\n", "})\n", "dataset = dataset.remove_columns('__index_level_0__')" ] }, { "cell_type": "code", "execution_count": null, "id": "1efb5845", "metadata": {}, "outputs": [], "source": [ "dataset.push_to_hub('yairschiff/zinc250k')" ] }, { "cell_type": "markdown", "id": "5c6f357d-20d9-4004-8091-68726b6b4c86", "metadata": {}, "source": [ "## Create tokenizer" ] }, { "cell_type": "code", "execution_count": 8, "id": "6642fc9d-4863-4e14-947b-95bae48e192d", "metadata": {}, "outputs": [], "source": [ "def smi_tokenizer(smi):\n", " \"\"\"Tokenize a SMILES molecule or reaction.\n", "\n", " Copied from https://github.com/pschwllr/MolecularTransformer.\n", " \"\"\"\n", " import re\n", " pattern = \"(\\[[^\\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\\(|\\)|\\.|=|#|-|\\+|\\\\\\\\|\\/|:|~|@|\\?|>|\\*|\\$|\\%[0-9]{2}|[0-9])\"\n", " regex = re.compile(pattern)\n", " tokens = [token for token in regex.findall(smi)]\n", " assert smi == ''.join(tokens)\n", " return tokens" ] }, { "cell_type": "code", "execution_count": 11, "id": "3a9e2e60-8596-4a91-acc3-d43e166ce723", "metadata": {}, "outputs": [], "source": [ "dataset = datasets.load_dataset('yairschiff/zinc250k')" ] }, { "cell_type": "code", "execution_count": 12, "id": "fbd5c2fe-4318-46bb-bc43-6ef7fe76e9fd", "metadata": {}, "outputs": [], "source": [ "# # If vocab file not created yet, uncomment and run this cell\n", "\n", "# tokens = []\n", "# for split in dataset.keys():\n", "# for smi in dataset[split]['canonical_smiles']:\n", "# tokens.extend(smi_tokenizer(smi))\n", "\n", "# with open('zinc250k_vocab.json', 'w', encoding='utf-8') as f:\n", "# f.write(\n", "# json.dumps(\n", "# {t: i for i, t in enumerate(sorted(set(tokens)))},\n", "# indent=2,\n", "# sort_keys=True,\n", "# ensure_ascii=False\n", "# ) + '\\n')" ] }, { "cell_type": "code", "execution_count": 14, "id": "4962478b-5343-4838-befe-64a5389625d4", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "CommitInfo(commit_url='https://huggingface.co/yairschiff/zinc250k-tokenizer/commit/7a07b0165a8a4f14f09d6137da8cdabf789397fd', commit_message='Upload tokenizer', commit_description='', oid='7a07b0165a8a4f14f09d6137da8cdabf789397fd', pr_url=None, pr_revision=None, pr_num=None)" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# # If HF tokenizer not yet published, uncomment and run this cell\n", "# import tokenizer\n", "\n", "# tokenizer.Zinc250kTokenizer.register_for_auto_class()\n", "# zinc250k_tokenizer = tokenizer.Zinc250kTokenizer(vocab_file='zinc250k_vocab.json')\n", "# zinc250k_tokenizer.push_to_hub('yairschiff/zinc250k-tokenizer')" ] }, { "cell_type": "code", "execution_count": 18, "id": "a779aa57-0c9d-4b8c-bf11-ccc5ab4c462e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Cn1ncc2c1CCC[C@H]2NC(=O)NC[C@H](O)COc1ccc(F)cc1\n", "[0, 25, 69, 15, 69, 68, 68, 16, 68, 15, 25, 25, 25, 35, 16, 29, 25, 11, 23, 30, 12, 29, 25, 35, 11, 30, 12, 25, 30, 68, 15, 68, 68, 68, 11, 27, 12, 68, 68, 15, 1]\n", "Cn1ncc2c1CCC[C@H]2NC(=O)NC[C@H](O)COc1ccc(F)cc1\n", "Cn1ncc2c1CCC[C@H]2NC(=O)NC[C@H](O)COc1ccc(F)cc1\n" ] } ], "source": [ "# Test tokenizer\n", "zinc250k_tokenizer = transformers.AutoTokenizer.from_pretrained(\n", " 'yairschiff/zinc250k-tokenizer', trust_remote_code=True, resume_download=None)\n", "print(dataset['train'][1000]['canonical_smiles'])\n", "print(zinc250k_tokenizer.encode(dataset['train'][1000]['canonical_smiles']))\n", "print(zinc250k_tokenizer.decode(zinc250k_tokenizer.encode(dataset['train'][1000]['canonical_smiles'])))\n", "print(zinc250k_tokenizer.decode(zinc250k_tokenizer.encode(dataset['train'][1000]['canonical_smiles'], add_special_tokens=False)))" ] }, { "cell_type": "code", "execution_count": 28, "id": "f3a15585-8e75-409d-9afe-0e7fe4a0bffc", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/224568 [00:00