{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":67356,"databundleVersionId":8006601,"sourceType":"competition"},{"sourceId":8042988,"sourceType":"datasetVersion","datasetId":4740586}],"dockerImageVersionId":30684,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# What is about ? \n\nNotebook gives an introductory example how to use ChemBerta to get embeddings for SMILES. \nIt is based on ALEKSEY TREPETSKY \"Create_ChemBERTa_embed\" notebook: https://www.kaggle.com/code/alekseytrepetsky/create-chemberta-embed\n(please upvote his notebook/dataset).  It is based on ChemBERTa-2 (not the ChemBERTa first version - as far as I understand)\n\nSee https://huggingface.co/DeepChem, https://deepchem.io/\n\nPapers:\n\n\"ChemBERTa-2: Towards Chemical Foundation Models\"  https://arxiv.org/abs/2209.01712\n\n\"ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction\" https://arxiv.org/abs/2010.09885\n\nTime estimates:\n\n    # CPU:  10 000 -  Wall time: 1min 47s\n    # GPU:  10 000 -  Wall time: 1min 16s - GPU should be much faster, may be the code below is not utilizing it correctly\n\nOverall we have 90 millions SMILES - so it would be too long to create embeds for all 90 millions SMILES on Kaggle by myself: either downsample, or crowdsource - split on many people : \n\n    10 000 - 2 minutes: \n    3 600 000 - 12 hours \n    So need about 30 Kaggle 12 hour runs, if we have 6 people we can do it in 12 hours since Kaggle allows to run 5 notebooks in parallel. \n    That is feasible, can be done later - if ChemBert will be really useful. \n\n\n---------------------------\n\nSome survey of different SMILEs embeds: https://medium.com/p/93246c75be6a  (by Sergey Kolchenko - Kaggle \"kvigly\" )\n\nAnother examples of SMILES embdeds  \"siamese-smole-bert\"  in SESHURAJUP notebook  https://www.kaggle.com/code/seshurajup/buildingblock-analysis-smiles-transformer see  https://huggingface.co/UdS-LSV/siamese-smole-bert-muv-1x by LSV @ Saarland University. That seems to be even more new than ChemBERTa-2 (since their repository links/thanks to ChemBERTa-2). \n\nPeople also recommend \"chemprop\": https://chemprop.readthedocs.io/en/latest/tutorial.html\n\n\n\n","metadata":{"execution":{"iopub.status.busy":"2024-04-12T20:22:16.476703Z","iopub.execute_input":"2024-04-12T20:22:16.477097Z","iopub.status.idle":"2024-04-12T20:22:16.481626Z","shell.execute_reply.started":"2024-04-12T20:22:16.477068Z","shell.execute_reply":"2024-04-12T20:22:16.480452Z"}}},{"cell_type":"markdown","source":"# Preliminaries","metadata":{}},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport time\nt0start = time.time() \n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-04-13T21:26:36.659846Z","iopub.execute_input":"2024-04-13T21:26:36.661081Z","iopub.status.idle":"2024-04-13T21:26:36.682053Z","shell.execute_reply.started":"2024-04-13T21:26:36.661046Z","shell.execute_reply":"2024-04-13T21:26:36.680974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pickle","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.683999Z","iopub.execute_input":"2024-04-13T21:26:36.684408Z","iopub.status.idle":"2024-04-13T21:26:36.689176Z","shell.execute_reply.started":"2024-04-13T21:26:36.684381Z","shell.execute_reply":"2024-04-13T21:26:36.688220Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load SMILES data","metadata":{}},{"cell_type":"markdown","source":"## Original non-economic format","metadata":{}},{"cell_type":"code","source":"%%time\n\ntrain = pd.read_csv('/kaggle/input/leash-BELKA/train.csv', nrows = 2)\ntrain.head()\n","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.690706Z","iopub.execute_input":"2024-04-13T21:26:36.691099Z","iopub.status.idle":"2024-04-13T21:26:36.747604Z","shell.execute_reply.started":"2024-04-13T21:26:36.691069Z","shell.execute_reply":"2024-04-13T21:26:36.746802Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Compact format by GREYSNOW\n\nPlease upvote GREYSNOW: \nhttps://www.kaggle.com/code/shlomoron/belka-shrunken-train-set-loading\n\n\nLoad full data - 5.5 minutes ","metadata":{}},{"cell_type":"code","source":"%%time\n\ndtypes = {'buildingblock1_smiles': np.int16, 'buildingblock2_smiles': np.int16, 'buildingblock3_smiles': np.int16,\n          'binds_BRD4':np.byte, 'binds_HSA':np.byte, 'binds_sEH':np.byte}\n\ntrain = pd.read_csv('/kaggle/input/belka-shrunken-train-set/train.csv', dtype = dtypes,  nrows = 10000 )\nprint(len(train))\ndisplay(train.head())\n","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.748784Z","iopub.execute_input":"2024-04-13T21:26:36.749665Z","iopub.status.idle":"2024-04-13T21:26:36.797313Z","shell.execute_reply.started":"2024-04-13T21:26:36.749635Z","shell.execute_reply":"2024-04-13T21:26:36.796593Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train.info()","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.799748Z","iopub.execute_input":"2024-04-13T21:26:36.800297Z","iopub.status.idle":"2024-04-13T21:26:36.829758Z","shell.execute_reply.started":"2024-04-13T21:26:36.800270Z","shell.execute_reply":"2024-04-13T21:26:36.828761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## building blocks data (not used in the current version of the notebook)","metadata":{}},{"cell_type":"code","source":"%%time\nBBs_dict_reverse_1 = pickle.load(open('/kaggle/input/belka-shrunken-train-set/train_dicts/BBs_dict_reverse_1.p', 'br'))\nBBs_dict_reverse_2 = pickle.load(open('/kaggle/input/belka-shrunken-train-set/train_dicts/BBs_dict_reverse_2.p', 'br'))\nBBs_dict_reverse_3 = pickle.load(open('/kaggle/input/belka-shrunken-train-set/train_dicts/BBs_dict_reverse_3.p', 'br'))","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.831064Z","iopub.execute_input":"2024-04-13T21:26:36.831473Z","iopub.status.idle":"2024-04-13T21:26:36.850338Z","shell.execute_reply.started":"2024-04-13T21:26:36.831446Z","shell.execute_reply":"2024-04-13T21:26:36.849398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nbuildingblock3_smiles_original = [BBs_dict_reverse_3[x] for x in train.buildingblock3_smiles[:1000]]\nprint(buildingblock3_smiles_original[0])","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.851586Z","iopub.execute_input":"2024-04-13T21:26:36.851869Z","iopub.status.idle":"2024-04-13T21:26:36.858547Z","shell.execute_reply.started":"2024-04-13T21:26:36.851846Z","shell.execute_reply":"2024-04-13T21:26:36.857385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# train.drop(['molecule_smiles'], axis = 1).to_parquet('train_lean.parquet', index = False)\n# pickle.dump(train.molecule_smiles.to_numpy(), open('molecule_smiles.p', 'bw'))","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.860289Z","iopub.execute_input":"2024-04-13T21:26:36.860727Z","iopub.status.idle":"2024-04-13T21:26:36.874054Z","shell.execute_reply.started":"2024-04-13T21:26:36.860689Z","shell.execute_reply":"2024-04-13T21:26:36.872970Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load model from Hugging Face \n\n\n### AutoModelForMaskedLM \n\nis a class that automatically selects the appropriate model architecture for masked language modeling (MLM) tasks. Masked language modeling is a type of language modeling where certain tokens in the input sequence are masked, and the model is trained to predict these masked tokens based on the surrounding context.\n\nFor example, you can use AutoModelForMaskedLM.from_pretrained() method to load a pre-trained model for masked language modeling, and Hugging Face will automatically select the correct model architecture based on the identifier or configuration you provide. This allows for easier experimentation and integration of different pre-trained models without needing to modify your code for each specific model architecture.\n \n\n###  Auto Tokenizer\n\n is a class designed to automatically select the appropriate tokenizer for a given pre-trained model. It works by identifying the model type or architecture from the provided model name or identifier string and instantiates the corresponding tokenizer.\n ","metadata":{}},{"cell_type":"code","source":"%%time\nimport torch\nfrom transformers import AutoModelForMaskedLM, AutoTokenizer\n\nchemberta = AutoModelForMaskedLM.from_pretrained(\"DeepChem/ChemBERTa-77M-MTR\")\ntokenizer = AutoTokenizer.from_pretrained(\"DeepChem/ChemBERTa-77M-MTR\")\n","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:36.875594Z","iopub.execute_input":"2024-04-13T21:26:36.876556Z","iopub.status.idle":"2024-04-13T21:26:46.660355Z","shell.execute_reply.started":"2024-04-13T21:26:36.876524Z","shell.execute_reply":"2024-04-13T21:26:46.659596Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Example with one SMILES\n\n    model_output = chemberta(**encoded_input)\n\nReturns essentially the matrix of vectors embedding for each token - modula some technical details.\n\nDetails are the following:\n\n    Take model_output[0] - tensor itself without technical things\n    It actually has 3 dimensions: [1, almost-number-tokens, embedding-dimension]  -  the first dimension is just 1 - so squeeze it\n    To get embdding for SMILES (not each token) we can either average of tokens dimension or take something like CLS - the zero layer\n    \n","metadata":{}},{"cell_type":"code","source":"%%time\nN = 1\nsmiles = train['molecule_smiles'].iat[0]#[:30]\nprint('Len of smiles string:',len(smiles))\nprint(smiles)\nwith torch.no_grad():\n    padding=True\n    encoded_input = tokenizer(smiles, return_tensors=\"pt\",padding=padding,truncation=True)\n    model_output = chemberta(**encoded_input)\n    print()\n    print('model_output:')\n    print('model_output info - type:', type(model_output),'len(model_output)', len(model_output) ,)\n    print('type( model_output[0]):', type( model_output[0]), 'model_output[0].shape:',  model_output[0].shape  )\n    print('model_output - first 100 symbols', str( model_output)[:100] )\n    print()\n    \n    embedding = model_output[0][:,0,:]\n    embeddings_cls = embedding\n    print('embeddings_cls: type', type(embeddings_cls), 'shape:', embeddings_cls.shape)\n    print()\n\n    embedding = torch.mean(model_output[0],1)\n    embeddings_mean = embedding\n    print('embeddings_mean: type', type(embeddings_mean), 'shape:', embeddings_mean.shape)\n    print()\n    \n# print(embeddings_cls)","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:46.661490Z","iopub.execute_input":"2024-04-13T21:26:46.662490Z","iopub.status.idle":"2024-04-13T21:26:46.792629Z","shell.execute_reply.started":"2024-04-13T21:26:46.662459Z","shell.execute_reply":"2024-04-13T21:26:46.791881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"embedding.shape","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:46.793917Z","iopub.execute_input":"2024-04-13T21:26:46.794450Z","iopub.status.idle":"2024-04-13T21:26:46.799684Z","shell.execute_reply.started":"2024-04-13T21:26:46.794421Z","shell.execute_reply":"2024-04-13T21:26:46.799027Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_output[0].shape","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:46.801264Z","iopub.execute_input":"2024-04-13T21:26:46.801577Z","iopub.status.idle":"2024-04-13T21:26:46.812469Z","shell.execute_reply.started":"2024-04-13T21:26:46.801552Z","shell.execute_reply":"2024-04-13T21:26:46.811465Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Aux Function for run through many SMILES","metadata":{}},{"cell_type":"code","source":"%%time\nfrom tqdm import tqdm\n\nchemberta.eval()\ndef featurize_ChemBERTa(smiles_list, padding=True):\n    embeddings_cls = torch.zeros(len(smiles_list), 600)\n    embeddings_mean = torch.zeros(len(smiles_list), 600)\n\n    with torch.no_grad():\n        for i, smiles in enumerate(tqdm(smiles_list)):\n            encoded_input = tokenizer(smiles, return_tensors=\"pt\",padding=padding,truncation=True)\n            model_output = chemberta(**encoded_input)\n            \n            embedding = model_output[0][::,0,::]\n            embeddings_cls[i] = embedding\n            \n            embedding = torch.mean(model_output[0],1)\n            embeddings_mean[i] = embedding\n            \n    return embeddings_cls.numpy(), embeddings_mean.numpy()\n","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:46.813636Z","iopub.execute_input":"2024-04-13T21:26:46.813950Z","iopub.status.idle":"2024-04-13T21:26:46.827231Z","shell.execute_reply.started":"2024-04-13T21:26:46.813920Z","shell.execute_reply":"2024-04-13T21:26:46.826263Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run through many SMILES","metadata":{}},{"cell_type":"code","source":"%%time\nN = int(1e4)\nsmiles_list = train['molecule_smiles'].iloc[:N].to_list()\nprint(smiles_list[:10])\ntrain_cls_pad_true, train_mean_pad_true = featurize_ChemBERTa(smiles_list)\n\n# CPU:  int(1e4) -  Wall time: 1min 47s\n# GPU:  int(1e4) -  Wall time: 1min 16s\n","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:26:46.832910Z","iopub.execute_input":"2024-04-13T21:26:46.833289Z","iopub.status.idle":"2024-04-13T21:28:22.523808Z","shell.execute_reply.started":"2024-04-13T21:26:46.833261Z","shell.execute_reply":"2024-04-13T21:28:22.523073Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_cls_pad_true.shape","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:28:22.525136Z","iopub.execute_input":"2024-04-13T21:28:22.525702Z","iopub.status.idle":"2024-04-13T21:28:22.532241Z","shell.execute_reply.started":"2024-04-13T21:28:22.525672Z","shell.execute_reply":"2024-04-13T21:28:22.531065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nnp.save('train_ChemBERTa_v2_77MTR_cls_pad_True.npy', train_cls_pad_true)\nnp.save('train_ChemBERTa_v2_77MTR_mean_pad_True.npy', train_mean_pad_true)\n","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:28:22.533582Z","iopub.execute_input":"2024-04-13T21:28:22.533948Z","iopub.status.idle":"2024-04-13T21:28:22.599160Z","shell.execute_reply.started":"2024-04-13T21:28:22.533891Z","shell.execute_reply":"2024-04-13T21:28:22.597996Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Clustermaps ","metadata":{}},{"cell_type":"code","source":"%%time\n\nfor emb in [train_cls_pad_true, train_mean_pad_true]:\n    \n    cm = np.corrcoef(emb[:100,:].T)\n    cm[np.isnan(cm)] = 0 \n    print(cm.shape)\n    sns.clustermap(cm, cmap='coolwarm')\n    plt.title('Correlation of emb coordinates',fontsize = 20)\n    plt.show()\n\n\n    cm = np.corrcoef(emb[:700,:])\n    cm[np.isnan(cm)] = 0 \n    print(cm.shape)    \n    sns.clustermap(cm, cmap='coolwarm')\n    plt.title('Correlations of smiles',fontsize = 20)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:30:20.874736Z","iopub.execute_input":"2024-04-13T21:30:20.875847Z","iopub.status.idle":"2024-04-13T21:30:28.908637Z","shell.execute_reply.started":"2024-04-13T21:30:20.875804Z","shell.execute_reply":"2024-04-13T21:30:28.906738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# UMAP","metadata":{}},{"cell_type":"code","source":"%%time\nimport umap\n\nreducer = umap.UMAP()\n\nfor emb in [train_cls_pad_true, train_mean_pad_true]:\n    r = reducer.fit_transform(emb)\n    sns.scatterplot(x = r[:,0],y = r[:,1])\n    plt.show()\n    d = pd.DataFrame(r)\n    d=d.reset_index()\n    display( d.corr() )\n\n    for col in [ 'binds_BRD4', 'binds_HSA', 'binds_sEH']:\n        sns.scatterplot(x = r[:,0],y = r[:,1], hue = train[col])\n        plt.title(col,fontsize = 20 )\n        plt.show()    ","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:37:37.863095Z","iopub.execute_input":"2024-04-13T21:37:37.863536Z","iopub.status.idle":"2024-04-13T21:38:07.545247Z","shell.execute_reply.started":"2024-04-13T21:37:37.863508Z","shell.execute_reply":"2024-04-13T21:38:07.543815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Final timing","metadata":{}},{"cell_type":"code","source":"print('%.1f seconds passed total '%(time.time()-t0start) )\nprint('%.1f minutes passed total '%( (time.time()-t0start)/60)  )\nprint('%.2f hours passed total '%( (time.time()-t0start)/3600)  )","metadata":{"execution":{"iopub.status.busy":"2024-04-13T21:28:23.979033Z","iopub.execute_input":"2024-04-13T21:28:23.979319Z","iopub.status.idle":"2024-04-13T21:28:23.984934Z","shell.execute_reply.started":"2024-04-13T21:28:23.979295Z","shell.execute_reply":"2024-04-13T21:28:23.984056Z"},"trusted":true},"execution_count":null,"outputs":[]}]}