{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":30461,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This notebook demonstrates how to use [ESM-2](https://github.com/facebookresearch/esm) transformer protein language model to extract embeddings from provided protein sequences. The following code calculates the embeddings for each protein sequence in the training and test FASTA files and saves them as individual `.pt` files. \n\nThe resulting embedding can be loaded using the following code:\n```\nimport torch\n\nembedding = torch.load('[EntryID].pt')\nembedding = embedding['mean_representations'][33].numpy()\n```\n\nComputing the embeddings and subsequently reading in the `.pt` files can take a while. The resulting numpy arrays can be found [here](https://www.kaggle.com/datasets/viktorfairuschin/cafa-5-ems-2-embeddings-numpy).\n\n**Note** that the test FASTA file contains duplicate entries. For this reason, this notebook uses cleaned FASTA files, which can be found [here](https://www.kaggle.com/datasets/viktorfairuschin/cafa-5-fasta-files).","metadata":{}},{"cell_type":"code","source":"!pip install -q fair-esm","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2025-11-01T10:13:44.941519Z","iopub.execute_input":"2025-11-01T10:13:44.941936Z","iopub.status.idle":"2025-11-01T10:13:53.197358Z","shell.execute_reply.started":"2025-11-01T10:13:44.941902Z","shell.execute_reply":"2025-11-01T10:13:53.196103Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm.auto import tqdm\nfrom esm import FastaBatchedDataset, pretrained\nimport joblib\nimport pathlib\nimport torch\nfrom tqdm.notebook import tqdm","metadata":{"execution":{"iopub.status.busy":"2025-11-01T10:16:57.924271Z","iopub.execute_input":"2025-11-01T10:16:57.925157Z","iopub.status.idle":"2025-11-01T10:17:08.836481Z","shell.execute_reply.started":"2025-11-01T10:16:57.925119Z","shell.execute_reply":"2025-11-01T10:17:08.835434Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_embeddings(model_name, fasta_file, output_dir, tokens_per_batch=4096, seq_length=1022, repr_layers=[33]):\n    import time\n    \n    model, alphabet = pretrained.load_model_and_alphabet(model_name)\n    model.eval()\n    if torch.cuda.is_available():\n        model = model.cuda()\n        \n    dataset = FastaBatchedDataset.from_file(fasta_file)\n    batches = dataset.get_batch_indices(tokens_per_batch, extra_toks_per_seq=1)\n    data_loader = torch.utils.data.DataLoader(\n        dataset, \n        collate_fn=alphabet.get_batch_converter(seq_length), \n        batch_sampler=batches\n    )\n    output_dir.mkdir(parents=True, exist_ok=True)\n    embedding_res = {}\n    \n    # Calculate total batches for progress tracking\n    total_batches = len(batches)\n    print(f\"Starting extraction: {total_batches} batches to process\")\n    print(\"-\" * 50)\n    \n    start_time = time.time()\n    \n    with torch.no_grad():\n        for batch_idx, (labels, strs, toks) in enumerate(data_loader):\n            # Progress tracking\n            if batch_idx % 100 == 0 or batch_idx == total_batches - 1:\n                elapsed_time = time.time() - start_time\n                progress = (batch_idx + 1) / total_batches\n                eta = (elapsed_time / (batch_idx + 1)) * (total_batches - batch_idx - 1)\n                \n                print(f\"Batch {batch_idx + 1}/{total_batches} | \"\n                      f\"Progress: {progress*100:.1f}% | \"\n                      f\"Elapsed: {elapsed_time:.1f}s | \"\n                      f\"ETA: {eta:.1f}s\")\n            \n            toks = toks.to(device=\"cuda\", non_blocking=True)\n            out = model(toks, repr_layers=repr_layers, return_contacts=False)\n            representations = {layer: t.to(device=\"cpu\") for layer, t in out[\"representations\"].items()}\n            \n            for i, label in enumerate(labels):\n                entry_id = label.split()[0]\n                truncate_len = min(seq_length, len(strs[i]))\n                \n                # Get the embedding from the last layer\n                final_layer = max(repr_layers)\n                embedding_res[entry_id] = representations[final_layer][i, 1 : truncate_len + 1].mean(0).clone().numpy()\n    \n    total_time = time.time() - start_time\n    print(\"-\" * 50)\n    print(f\"✓ Extraction complete!\")\n    print(f\"  Total time: {total_time:.2f} seconds\")\n    print(f\"  Proteins processed: {len(embedding_res)}\")\n    print(f\"  Average time per batch: {total_time/total_batches:.2f}s\")\n    \n    joblib.dump(embedding_res, 'embedding_test.joblib')\n    print(f\"✓ Saved embeddings to embedding_test.joblib\")\n    \n\nmodel_name = 'esm2_t33_650M_UR50D'\nfasta_file = pathlib.Path('/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta')\noutput_dir = pathlib.Path('test_embeddings')\nextract_embeddings(model_name, fasta_file, output_dir)","metadata":{"execution":{"iopub.status.busy":"2025-11-01T10:18:28.414038Z","iopub.execute_input":"2025-11-01T10:18:28.414995Z","iopub.status.idle":"2025-11-01T10:19:27.654171Z","shell.execute_reply.started":"2025-11-01T10:18:28.414957Z","shell.execute_reply":"2025-11-01T10:19:27.652673Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Process train file","metadata":{}},{"cell_type":"markdown","source":"## Process test file","metadata":{}},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}