{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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)\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","trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:24:05.921307Z","iopub.execute_input":"2026-01-26T22:24:05.922002Z","iopub.status.idle":"2026-01-26T22:24:05.931322Z","shell.execute_reply.started":"2026-01-26T22:24:05.921973Z","shell.execute_reply":"2026-01-26T22:24:05.930722Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Biopython --q\n!pip install transformers accelerate sentencepiece biopython --quiet\n!pip install fair-esm --q\n!pip install tqdm\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:24:05.932876Z","iopub.execute_input":"2026-01-26T22:24:05.933120Z","iopub.status.idle":"2026-01-26T22:24:22.055889Z","shell.execute_reply.started":"2026-01-26T22:24:05.933098Z","shell.execute_reply":"2026-01-26T22:24:22.054950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load protein sequences\nfrom Bio import SeqIO\nimport pandas as pd\ntrain_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\ntest_path = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\ntrain_records = list(SeqIO.parse(train_path,\"fasta\"))\ndf_train = pd.DataFrame({\n    \"protein_id\":[rec.id.split(\"|\")[1] for rec in train_records],\n    \"sequence\":[str(rec.seq) for rec in train_records],\n})\n\ntest_records = list(SeqIO.parse(test_path,\"fasta\"))\ndf_test = pd.DataFrame({\n    \"protein_id\":[rec.id.split(\" \")[0] for rec in test_records],\n    \"sequence\":[str(rec.seq) for rec in test_records],\n})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:24:22.057228Z","iopub.execute_input":"2026-01-26T22:24:22.057560Z","iopub.status.idle":"2026-01-26T22:24:26.471491Z","shell.execute_reply.started":"2026-01-26T22:24:22.057501Z","shell.execute_reply":"2026-01-26T22:24:26.470900Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Use  ESM2‑T30‑150M \n","metadata":{}},{"cell_type":"code","source":"train_data = list(zip(df_train[\"protein_id\"], df_train[\"sequence\"]))\ntest_data = list(zip(df_test[\"protein_id\"], df_test[\"sequence\"]))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:24:26.472374Z","iopub.execute_input":"2026-01-26T22:24:26.472677Z","iopub.status.idle":"2026-01-26T22:24:26.611710Z","shell.execute_reply.started":"2026-01-26T22:24:26.472640Z","shell.execute_reply":"2026-01-26T22:24:26.610864Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport esm\nimport os\nfrom tqdm import tqdm\nimport gc\n\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ntorch.set_grad_enabled(False)\n\nmodel_name = 'esm2_t30_150M_UR50D'\nmodel, alphabet = esm.pretrained.esm2_t30_150M_UR50D()\nmodel = model.half().to(device)\nmodel.eval()\nbatch_converter = alphabet.get_batch_converter()\n\nmodel_layers = {\n    'esm2_t6_8M_UR50D': 6,\n    'esm2_t12_35M_UR50D': 12,\n    'esm2_t30_150M_UR50D': 30,\n    'esm2_t33_650M_UR50D': 33,\n    'esm2_t36_3B_UR50D': 36,\n}\n\nlast_layer = model_layers[model_name]\n\ndef extract_esm2_embeddings(data, batch_size=1, max_seq_len=1022, \n                            save_every=None, save_path=None):\n    \"\"\"\n    Extract ESM2 embeddings with OOM protection\n    \n    Args:\n        save_every: Save checkpoint every N batches (None to disable)\n        save_path: Base path for checkpoints\n    \"\"\"\n    embeddings = {}\n    \n    # Sort by sequence length for efficient batching\n    data_sorted = sorted(data, key=lambda x: len(x[1]))\n    \n    # Group by similar lengths\n    batches = []\n    current_batch = []\n    current_length = 0\n    \n    for protein_id, seq in data_sorted:\n        if len(seq) > max_seq_len:\n            print(f\"Warning: Truncating {protein_id} from {len(seq)} to {max_seq_len}\")\n            seq = seq[:max_seq_len]\n        \n        if current_batch and (len(seq) > current_length * 1.5 or len(current_batch) >= batch_size):\n            batches.append(current_batch)\n            current_batch = []\n            current_length = 0\n        \n        current_batch.append((protein_id, seq))\n        current_length = max(current_length, len(seq))\n    \n    if current_batch:\n        batches.append(current_batch)\n    \n    print(f\"Processing {len(data)} proteins in {len(batches)} batches\")\n    \n    for batch_idx, batch in enumerate(tqdm(batches, desc=\"Extracting ESM2 embeddings\")):\n        try:\n            labels, strs, toks = batch_converter(batch)\n            toks = toks.to(device)\n            \n            with torch.no_grad():\n                out = model(toks, repr_layers=[last_layer], return_contacts=False)\n            \n            reps = out[\"representations\"][last_layer]\n            \n            for idx, (protein_id, seq) in enumerate(batch):\n                rep = reps[idx, 1:len(seq)+1]\n                emb = rep.mean(0).cpu().numpy().astype(np.float32)\n                embeddings[protein_id] = emb\n            \n            del toks, out, reps\n            torch.cuda.empty_cache()\n            \n            # Only save if save_every is specified AND we've hit the interval\n            if save_every and save_path and (batch_idx + 1) % save_every == 0:\n                checkpoint_file = f\"{save_path}_checkpoint_{batch_idx}.npy\"\n                np.save(checkpoint_file, embeddings)\n                print(f\"Saved checkpoint at batch {batch_idx} ({len(embeddings)} proteins)\")\n        \n        except RuntimeError as e:\n            if \"out of memory\" in str(e):\n                print(f\"OOM at batch {batch_idx}, skipping\")\n                torch.cuda.empty_cache()\n                gc.collect()\n                continue\n            else:\n                raise e\n    \n    return embeddings\n\n\n# USAGE:\n\n# Training set - no checkpoints needed (small dataset)\nprint(\"Processing training data...\")\ntrain_embeddings = extract_esm2_embeddings(\n    train_data,\n    batch_size=4,\n    save_every=None,  # No checkpoints\n    save_path=None\n)\n\n# Save final training embeddings\nnp.save(\"/kaggle/working/train_esm2_embeddings.npy\", train_embeddings)\nprint(f\"Saved {len(train_embeddings)} training embeddings\")\n\n# Clear memory before test set\ndel train_embeddings\ngc.collect()\ntorch.cuda.empty_cache()\n\n# Test set - process in chunks and save each chunk immediately\nprint(\"\\nProcessing test data in chunks...\")\nchunk_size = 5000  # Smaller chunks to manage disk space\n\nfor chunk_idx in range(0, len(test_data), chunk_size):\n    chunk_num = chunk_idx // chunk_size + 1\n    print(f\"\\n=== Chunk {chunk_num}/{(len(test_data)-1)//chunk_size + 1} ===\")\n    \n    chunk = test_data[chunk_idx:chunk_idx + chunk_size]\n    \n    chunk_embeddings = extract_esm2_embeddings(\n        chunk,\n        batch_size=2,\n        save_every=None,  # Don't save intermediate checkpoints\n        save_path=None\n    )\n    \n    # Save this chunk immediately\n    chunk_file = f\"/kaggle/working/test_esm2_chunk_{chunk_num}.npy\"\n    np.save(chunk_file, chunk_embeddings)\n    print(f\"Saved chunk {chunk_num} with {len(chunk_embeddings)} proteins\")\n    \n    # Clear memory\n    del chunk_embeddings, chunk\n    gc.collect()\n    torch.cuda.empty_cache()\n\nprint(\"\\n✅ All embeddings generated and saved!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T22:24:26.613481Z","iopub.execute_input":"2026-01-26T22:24:26.613782Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_embeddings = extract_esm2_embeddings(test_data, batch_size=1)\nX_test = np.vstack([test_embeddings[pid] for pid in df_test[\"protein_id\"]])\nnp.save(\"X_test_esm2.npy\", X_test)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}