{"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"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":41875,"databundleVersionId":5521661,"sourceType":"competition"}],"dockerImageVersionId":30513,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install fair-esm","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:05:30.989502Z","iopub.execute_input":"2023-06-20T09:05:30.990007Z","iopub.status.idle":"2023-06-20T09:05:48.117705Z","shell.execute_reply.started":"2023-06-20T09:05:30.989973Z","shell.execute_reply":"2023-06-20T09:05:48.116475Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio import SeqIO\nimport numpy as np\nimport torch\nimport esm\nimport pandas as pd","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:05:48.123359Z","iopub.execute_input":"2023-06-20T09:05:48.125700Z","iopub.status.idle":"2023-06-20T09:05:51.549163Z","shell.execute_reply.started":"2023-06-20T09:05:48.125655Z","shell.execute_reply":"2023-06-20T09:05:51.548234Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load train FASTA","metadata":{}},{"cell_type":"code","source":"file_fasta = \"/kaggle/input/cafa-5-protein-function-prediction/Train/train_sequences.fasta\"\nsequences = SeqIO.parse(file_fasta, \"fasta\")\nlist_ids = [seq.id for seq in sequences]\nsequences = SeqIO.parse(file_fasta, \"fasta\")\nlist_lens = [len(seq.seq) for seq in sequences]\nsequences = SeqIO.parse(file_fasta, \"fasta\")\nlist_seqs = [str(seq.seq) for seq in sequences]","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:05:51.551564Z","iopub.execute_input":"2023-06-20T09:05:51.552532Z","iopub.status.idle":"2023-06-20T09:05:58.762182Z","shell.execute_reply.started":"2023-06-20T09:05:51.552492Z","shell.execute_reply":"2023-06-20T09:05:58.761127Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_seqs = pd.DataFrame(data = [list_ids, list_lens, list_seqs]).T\ndf_seqs.columns = ['id','len','seq']\ndf_seqs = df_seqs[df_seqs[\"len\"] > 1022]","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:05:58.764756Z","iopub.execute_input":"2023-06-20T09:05:58.765422Z","iopub.status.idle":"2023-06-20T09:06:05.523532Z","shell.execute_reply.started":"2023-06-20T09:05:58.765364Z","shell.execute_reply":"2023-06-20T09:06:05.522525Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(df_seqs.shape)\ndf_seqs.head(3)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:06:05.525066Z","iopub.execute_input":"2023-06-20T09:06:05.525780Z","iopub.status.idle":"2023-06-20T09:06:05.547037Z","shell.execute_reply.started":"2023-06-20T09:06:05.525746Z","shell.execute_reply":"2023-06-20T09:06:05.546166Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"list_ids = df_seqs[\"id\"].values.tolist()\nlist_seqs = df_seqs[\"seq\"].values.tolist()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:06:05.551616Z","iopub.execute_input":"2023-06-20T09:06:05.553865Z","iopub.status.idle":"2023-06-20T09:06:05.582128Z","shell.execute_reply.started":"2023-06-20T09:06:05.553833Z","shell.execute_reply":"2023-06-20T09:06:05.581140Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load model","metadata":{}},{"cell_type":"code","source":"# Load ESM-2 model\n# model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()\nmodel, alphabet = esm.pretrained.esm2_t36_3B_UR50D()\nbatch_converter = alphabet.get_batch_converter()\nmodel.eval()  # disables dropout for deterministic results\n\nmodel = model.cuda() # uncomment to run on GPU","metadata":{"execution":{"iopub.status.busy":"2023-06-20T06:31:51.731121Z","iopub.execute_input":"2023-06-20T06:31:51.732219Z","iopub.status.idle":"2023-06-20T06:32:19.003540Z","shell.execute_reply.started":"2023-06-20T06:31:51.732179Z","shell.execute_reply":"2023-06-20T06:32:19.002499Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Check with single sequence","metadata":{}},{"cell_type":"code","source":"# embed_layer = 33\nembed_layer = 36\n\nmax_allowed_len = 1022\n\n# List for all protein embeds (mean of residual embeds)\nprotein_embeds = []\nresidual_embeds = []\n\nid = list_ids[0]\n# seq = list_seqs[0]\nseq = df_seqs[df_seqs[\"id\"] == \"O48653\"][\"seq\"].values.tolist()\nseq = seq[0]\nlen(seq)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:46:56.143385Z","iopub.execute_input":"2023-06-20T07:46:56.143749Z","iopub.status.idle":"2023-06-20T07:46:56.154760Z","shell.execute_reply.started":"2023-06-20T07:46:56.143721Z","shell.execute_reply":"2023-06-20T07:46:56.153882Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# List for all residual embeds for current protein\nresidual_embeds = []\n    \n# Embeds of first 1022 residuals\ndata = [(id, seq[0:min(len(seq), max_allowed_len)])]\nbatch_labels, batch_strs, batch_tokens = batch_converter(data)\nbatch_tokens = batch_tokens.cuda() \ntokens_len = (batch_tokens != alphabet.padding_idx).sum(1)\nwith torch.no_grad():\n    results = model(batch_tokens, repr_layers=[embed_layer], return_contacts=False)\ntoken_representations = results[\"representations\"][embed_layer]\ntoken_representations = token_representations[:, 1:(tokens_len-1), :]\ntoken_representations = token_representations.squeeze(dim = 0)\nresidual_embeds.append(token_representations)\n    \n# Remove first 511 residuals\nseq = seq[max_allowed_len//2:]\n    \nwhile len(seq) > 0:\n    data = [(id, seq[0:min(len(seq), max_allowed_len)])]\n    batch_labels, batch_strs, batch_tokens = batch_converter(data)\n    batch_tokens = batch_tokens.cuda() \n    tokens_len = (batch_tokens != alphabet.padding_idx).sum(1)\n    with torch.no_grad():\n        results = model(batch_tokens, repr_layers=[embed_layer], return_contacts=False)\n    token_representations = results[\"representations\"][embed_layer]\n    token_representations = token_representations[:, 1:(tokens_len-1), :]\n    token_representations = token_representations[:, max_allowed_len//2:, :]\n    token_representations = token_representations.squeeze(dim = 0)\n    residual_embeds.append(token_representations)\n                 \n    seq = seq[max_allowed_len//2:]","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:46:57.337391Z","iopub.execute_input":"2023-06-20T07:46:57.337857Z","iopub.status.idle":"2023-06-20T07:46:58.172355Z","shell.execute_reply.started":"2023-06-20T07:46:57.337795Z","shell.execute_reply":"2023-06-20T07:46:58.171399Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"residual_embeds[2].shape","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:47:05.811515Z","iopub.execute_input":"2023-06-20T07:47:05.811981Z","iopub.status.idle":"2023-06-20T07:47:05.823496Z","shell.execute_reply.started":"2023-06-20T07:47:05.811942Z","shell.execute_reply":"2023-06-20T07:47:05.822424Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"residual_embeds = torch.cat(residual_embeds)\nresidual_embeds.shape ","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:47:11.033027Z","iopub.execute_input":"2023-06-20T07:47:11.033471Z","iopub.status.idle":"2023-06-20T07:47:11.041890Z","shell.execute_reply.started":"2023-06-20T07:47:11.033420Z","shell.execute_reply":"2023-06-20T07:47:11.041012Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"residual_embeds = residual_embeds.mean(0)\nresidual_embeds.shape","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:20:10.426167Z","iopub.execute_input":"2023-06-20T07:20:10.426599Z","iopub.status.idle":"2023-06-20T07:20:10.434637Z","shell.execute_reply.started":"2023-06-20T07:20:10.426559Z","shell.execute_reply":"2023-06-20T07:20:10.433674Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":" residual_embeds = residual_embeds.cpu().numpy()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:20:14.756726Z","iopub.execute_input":"2023-06-20T07:20:14.757209Z","iopub.status.idle":"2023-06-20T07:20:14.762311Z","shell.execute_reply.started":"2023-06-20T07:20:14.757173Z","shell.execute_reply":"2023-06-20T07:20:14.761187Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Embeds for all train sequences > 1022\n","metadata":{}},{"cell_type":"code","source":"!mkdir train_embeds\n\n# embed_layer = 33\nembed_layer = 36\n\nmax_allowed_len = 1022\n\n# List for all protein embeds (mean of residual embeds)\nprotein_embeds = []\n    \n\nfor i, (id, seq) in enumerate(zip(list_ids, list_seqs)):\n    \n    if i % 100 == 0:\n        print(i,'sequences processed')\n    \n    # List for all residual embeds for current protein\n    residual_embeds = []\n    \n    # Embeds of first 1022 residuals\n    data = [(id, seq[0:min(len(seq), max_allowed_len)])]\n    batch_labels, batch_strs, batch_tokens = batch_converter(data)\n    batch_tokens = batch_tokens.cuda() \n    tokens_len = (batch_tokens != alphabet.padding_idx).sum(1)\n    with torch.no_grad():\n        results = model(batch_tokens, repr_layers=[embed_layer], return_contacts=False)\n    token_representations = results[\"representations\"][embed_layer]\n    token_representations = token_representations[:, 1:(tokens_len-1), :]\n    token_representations = token_representations.squeeze(dim=0)\n    residual_embeds.append(token_representations)\n    \n    # Remove first 511 residuals\n    seq = seq[max_allowed_len//2:]\n    \n    while len(seq) > 0:\n        data = [(id, seq[0:min(len(seq), max_allowed_len)])]\n        batch_labels, batch_strs, batch_tokens = batch_converter(data)\n        batch_tokens = batch_tokens.cuda() \n        tokens_len = (batch_tokens != alphabet.padding_idx).sum(1)\n        with torch.no_grad():\n            results = model(batch_tokens, repr_layers=[embed_layer], return_contacts=False)\n        token_representations = results[\"representations\"][embed_layer]\n        token_representations = token_representations[:, 1:(tokens_len-1), :]\n        # All but first 511 residuals\n        token_representations = token_representations[:, max_allowed_len//2:, :]\n        token_representations = token_representations.squeeze(dim=0)\n        residual_embeds.append(token_representations)\n                 \n        seq = seq[max_allowed_len//2:]\n    \n                 \n    residual_embeds = torch.cat(residual_embeds)\n    residual_embeds = residual_embeds.mean(0)\n    residual_embeds = residual_embeds.cpu().numpy()\n    np.save(f\"train_embeds/{id}.npy\", residual_embeds)           \n    # protein_embeds.append(residual_embeds)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T07:57:14.285510Z","iopub.execute_input":"2023-06-20T07:57:14.285926Z","iopub.status.idle":"2023-06-20T08:11:15.713299Z","shell.execute_reply.started":"2023-06-20T07:57:14.285894Z","shell.execute_reply":"2023-06-20T08:11:15.712235Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# protein_embeds = np.stack(protein_embeds)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T08:11:24.308526Z","iopub.execute_input":"2023-06-20T08:11:24.308923Z","iopub.status.idle":"2023-06-20T08:11:24.319178Z","shell.execute_reply.started":"2023-06-20T08:11:24.308890Z","shell.execute_reply":"2023-06-20T08:11:24.318222Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# protein_embeds.shape","metadata":{"execution":{"iopub.status.busy":"2023-06-20T08:11:26.028009Z","iopub.execute_input":"2023-06-20T08:11:26.028448Z","iopub.status.idle":"2023-06-20T08:11:26.035224Z","shell.execute_reply.started":"2023-06-20T08:11:26.028410Z","shell.execute_reply":"2023-06-20T08:11:26.034367Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# np.save(\"train_full_embeds.npy\", protein_embeds)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T08:11:35.407202Z","iopub.execute_input":"2023-06-20T08:11:35.407580Z","iopub.status.idle":"2023-06-20T08:11:35.417941Z","shell.execute_reply.started":"2023-06-20T08:11:35.407550Z","shell.execute_reply":"2023-06-20T08:11:35.416960Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_seqs.to_csv(\"train_ids.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-06-20T09:06:44.010533Z","iopub.execute_input":"2023-06-20T09:06:44.010991Z","iopub.status.idle":"2023-06-20T09:06:45.485467Z","shell.execute_reply.started":"2023-06-20T09:06:44.010956Z","shell.execute_reply":"2023-06-20T09:06:45.484366Z"},"trusted":true},"outputs":[],"execution_count":null}]}