{"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-10T09:15:05.403106Z","iopub.execute_input":"2026-01-10T09:15:05.403394Z","iopub.status.idle":"2026-01-10T09:15:06.445298Z","shell.execute_reply.started":"2026-01-10T09:15:05.403364Z","shell.execute_reply":"2026-01-10T09:15:06.444526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Biopython --q\n!pip install goatools --q","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:15:09.286384Z","iopub.execute_input":"2026-01-10T09:15:09.286665Z","iopub.status.idle":"2026-01-10T09:15:47.729988Z","shell.execute_reply.started":"2026-01-10T09:15:09.286641Z","shell.execute_reply":"2026-01-10T09:15:47.729285Z"}},"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-10T09:15:47.731531Z","iopub.execute_input":"2026-01-10T09:15:47.731784Z","iopub.status.idle":"2026-01-10T09:15:51.771763Z","shell.execute_reply.started":"2026-01-10T09:15:47.731756Z","shell.execute_reply":"2026-01-10T09:15:51.77121Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load train_taxonomy\ntaxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\"\ndf_tax = pd.read_csv(taxonomy_path,sep=\"\\t\", header=None)\ndf_tax.columns = [\"protein_id\", \"taxon_id\"]\ndf_tax.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:01.942725Z","iopub.execute_input":"2026-01-10T09:16:01.943014Z","iopub.status.idle":"2026-01-10T09:16:02.012983Z","shell.execute_reply.started":"2026-01-10T09:16:01.942987Z","shell.execute_reply":"2026-01-10T09:16:02.012408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n\ndf_terms = pd.read_csv(train_terms_path, sep=\"\\t\")\n\ndf_terms.columns = [\"protein_id\", \"go_id\", \"aspect\"]\ndf_terms.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:05.389293Z","iopub.execute_input":"2026-01-10T09:16:05.389771Z","iopub.status.idle":"2026-01-10T09:16:05.657487Z","shell.execute_reply.started":"2026-01-10T09:16:05.389742Z","shell.execute_reply":"2026-01-10T09:16:05.656914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Combine everything into a unified table\ndf_train = df_train.merge(df_tax, on = \"protein_id\",how=\"left\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:10.111359Z","iopub.execute_input":"2026-01-10T09:16:10.112018Z","iopub.status.idle":"2026-01-10T09:16:10.165005Z","shell.execute_reply.started":"2026-01-10T09:16:10.111977Z","shell.execute_reply":"2026-01-10T09:16:10.16429Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"terms_per_protein = (\n    df_terms.groupby(\"protein_id\")[\"go_id\"]\n    .apply(list)\n    .to_dict()\n)\n\ndf_train[\"go_terms\"] = df_train[\"protein_id\"].map(terms_per_protein).fillna(\"\").apply(\n    lambda x:x if isinstance(x,list) else[]\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:14.159512Z","iopub.execute_input":"2026-01-10T09:16:14.160059Z","iopub.status.idle":"2026-01-10T09:16:15.426028Z","shell.execute_reply.started":"2026-01-10T09:16:14.160031Z","shell.execute_reply":"2026-01-10T09:16:15.425425Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Proteins:\", len(df_train))\nprint(\"Proteins with GO terms:\", df_train['go_terms'].apply(len).gt(0).sum())\nprint(\"Unique GO terms:\", df_terms['go_id'].nunique())\nprint(\"Aspect distribution:\", df_terms['aspect'].value_counts())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:17.996979Z","iopub.execute_input":"2026-01-10T09:16:17.99737Z","iopub.status.idle":"2026-01-10T09:16:18.096826Z","shell.execute_reply.started":"2026-01-10T09:16:17.997341Z","shell.execute_reply":"2026-01-10T09:16:18.096223Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load test dataset\ntest_records = list(SeqIO.parse(test_path, \"fasta\"))\n\nprotein_ids = []\ntaxon_ids = []\nsequences = []\n\nfor rec in test_records:\n    header = rec.description.split()\n    protein_id = header[0]          # A0A0C5B5G6\n    taxon_id = header[1]            # 9606\n    \n    protein_ids.append(protein_id)\n    taxon_ids.append(taxon_id)\n    sequences.append(str(rec.seq))\n\ndf_test = pd.DataFrame({\n    \"protein_id\": protein_ids,\n    \"taxon_id\": taxon_ids,\n    \"sequence\": sequences\n})\n\ndf_test.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:21.037568Z","iopub.execute_input":"2026-01-10T09:16:21.037846Z","iopub.status.idle":"2026-01-10T09:16:22.932995Z","shell.execute_reply.started":"2026-01-10T09:16:21.037822Z","shell.execute_reply":"2026-01-10T09:16:22.932405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_tax_species=pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\", sep=\"\\t\")\ndf_tax_species.columns=[\"taxon_id\",\"species\"]\ndf_tax_species[\"taxon_id\"] = df_tax_species[\"taxon_id\"].astype(str)\ndf_tax_species.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:28.791455Z","iopub.execute_input":"2026-01-10T09:16:28.792337Z","iopub.status.idle":"2026-01-10T09:16:28.816459Z","shell.execute_reply.started":"2026-01-10T09:16:28.792298Z","shell.execute_reply":"2026-01-10T09:16:28.815862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_test = df_test.merge(df_tax_species, on=\"taxon_id\", how=\"left\")\ndf_test.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:33.037874Z","iopub.execute_input":"2026-01-10T09:16:33.038179Z","iopub.status.idle":"2026-01-10T09:16:33.096765Z","shell.execute_reply.started":"2026-01-10T09:16:33.038129Z","shell.execute_reply":"2026-01-10T09:16:33.096217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# GO label propagation\n\nfrom goatools.obo_parser import GODag\ngo_dag = GODag(\"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\")\n\n# Helper - get ancestors for a GO term\ndef get_ancestors(go_id,go_dag):\n    if go_id not in go_dag:\n        return set()\n    term = go_dag[go_id]\n    ancestors = term.get_all_parents()\n    ancestors.add(go_id)\n    return ancestors\n\n# Propagate GO terms for one protein\n\ndef propagate_terms(go_list, go_dag):\n    \"\"\"Propagate a list of GO terms to include all ancestors.\"\"\"\n    propagated = set()\n    for go_id in go_list:\n        if go_id not in go_dag:\n            continue\n        term = go_dag[go_id]\n        propagated.add(go_id)\n        ancestors = term.get_all_parents()\n        propagated |= ancestors\n    return propagated\n\ndf_train[\"go_terms_propagated\"]=df_train[\"go_terms\"].apply(\n    lambda terms: propagate_terms(terms,go_dag)\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:37.478811Z","iopub.execute_input":"2026-01-10T09:16:37.479507Z","iopub.status.idle":"2026-01-10T09:16:43.274054Z","shell.execute_reply.started":"2026-01-10T09:16:37.479477Z","shell.execute_reply":"2026-01-10T09:16:43.273231Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:43.275349Z","iopub.execute_input":"2026-01-10T09:16:43.275893Z","iopub.status.idle":"2026-01-10T09:16:43.285727Z","shell.execute_reply.started":"2026-01-10T09:16:43.275869Z","shell.execute_reply":"2026-01-10T09:16:43.285207Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Split MF / BP / C\ndef split_by_namespace(go_set, go_dag):\n    mf, bp, cc = [], [], []\n    for go_id in go_set:\n        if go_id in go_dag:\n            ns = go_dag[go_id].namespace\n            if ns == \"molecular_function\":\n                mf.append(go_id)\n            elif ns == \"biological_process\":\n                bp.append(go_id)\n            elif ns == \"cellular_component\":\n                cc.append(go_id)\n    return mf, bp, cc\n\ndf_train[\"MF\"], df_train[\"BP\"], df_train[\"CC\"] = zip(*df_train[\"go_terms_propagated\"].apply(\n    lambda s: split_by_namespace(s, go_dag)\n))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:51.744692Z","iopub.execute_input":"2026-01-10T09:16:51.744948Z","iopub.status.idle":"2026-01-10T09:16:53.2255Z","shell.execute_reply.started":"2026-01-10T09:16:51.744927Z","shell.execute_reply":"2026-01-10T09:16:53.224695Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:53.226678Z","iopub.execute_input":"2026-01-10T09:16:53.226944Z","iopub.status.idle":"2026-01-10T09:16:53.241062Z","shell.execute_reply.started":"2026-01-10T09:16:53.226922Z","shell.execute_reply":"2026-01-10T09:16:53.240402Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.preprocessing import MultiLabelBinarizer\nmlb_mf = MultiLabelBinarizer()\nY_mf = mlb_mf.fit_transform(df_train[\"MF\"])\n\nmlb_bp = MultiLabelBinarizer()\nY_bp = mlb_bp.fit_transform(df_train[\"BP\"])\n\nmlb_cc = MultiLabelBinarizer()\nY_cc = mlb_cc.fit_transform(df_train[\"CC\"])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:16:56.616332Z","iopub.execute_input":"2026-01-10T09:16:56.616912Z","iopub.status.idle":"2026-01-10T09:17:00.810446Z","shell.execute_reply.started":"2026-01-10T09:16:56.616886Z","shell.execute_reply":"2026-01-10T09:17:00.80991Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Y_mf.shape:\",Y_mf.shape)\nprint(\"Y_bp.shape:\",Y_bp.shape)\nprint(\"Y_cc.shape:\",Y_cc.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:17:05.342996Z","iopub.execute_input":"2026-01-10T09:17:05.34377Z","iopub.status.idle":"2026-01-10T09:17:05.347927Z","shell.execute_reply.started":"2026-01-10T09:17:05.343739Z","shell.execute_reply":"2026-01-10T09:17:05.347213Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# filter rare GO terms\n\nmin_count = 30  # common threshold\n\nmf_counts = Y_mf.sum(axis=0)\nkeep_mf = mf_counts >= min_count\nY_mf = Y_mf[:, keep_mf]\nmf_terms = [t for t, k in zip(mlb_mf.classes_, keep_mf) if k]\n\n\nmin_count = 50  # common threshold\nbp_counts = Y_bp.sum(axis=0)\nkeep_bp = bp_counts >= min_count\nY_bp = Y_bp[:, keep_bp]\nbp_terms = [t for t, k in zip(mlb_bp.classes_, keep_bp) if k]\n\nmin_count = 20\ncc_counts = Y_cc.sum(axis=0)\nkeep_cc = cc_counts >= min_count\nY_cc = Y_cc[:, keep_cc]\ncc_terms = [t for t, k in zip(mlb_bp.classes_, keep_bp) if k]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:17:11.518676Z","iopub.execute_input":"2026-01-10T09:17:11.5193Z","iopub.status.idle":"2026-01-10T09:17:25.575605Z","shell.execute_reply.started":"2026-01-10T09:17:11.51927Z","shell.execute_reply":"2026-01-10T09:17:25.574887Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Y_mf.shape:\",Y_mf.shape)\nprint(\"Y_bp.shape:\",Y_bp.shape)\nprint(\"Y_cc.shape:\",Y_cc.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:17:32.870371Z","iopub.execute_input":"2026-01-10T09:17:32.870636Z","iopub.status.idle":"2026-01-10T09:17:32.874943Z","shell.execute_reply.started":"2026-01-10T09:17:32.870615Z","shell.execute_reply":"2026-01-10T09:17:32.874372Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import joblib\n\njoblib.dump(mlb_mf, \"/kaggle/working/mlb_mf.pkl\")\njoblib.dump(mlb_bp, \"/kaggle/working/mlb_bp.pkl\")\njoblib.dump(mlb_cc, \"/kaggle/working/mlb_cc.pkl\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:17:34.014677Z","iopub.execute_input":"2026-01-10T09:17:34.014958Z","iopub.status.idle":"2026-01-10T09:17:34.027716Z","shell.execute_reply.started":"2026-01-10T09:17:34.014934Z","shell.execute_reply":"2026-01-10T09:17:34.027017Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Use  ESM2‑T33‑650M \n","metadata":{}},{"cell_type":"code","source":"!pip install fair-esm --q\n!pip install tqdm\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:17:53.519611Z","iopub.execute_input":"2026-01-10T09:17:53.519891Z","iopub.status.idle":"2026-01-10T09:17:59.899015Z","shell.execute_reply.started":"2026-01-10T09:17:53.519867Z","shell.execute_reply":"2026-01-10T09:17:59.898357Z"}},"outputs":[],"execution_count":null},{"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-10T09:18:03.295002Z","iopub.execute_input":"2026-01-10T09:18:03.295615Z","iopub.status.idle":"2026-01-10T09:18:03.506116Z","shell.execute_reply.started":"2026-01-10T09:18:03.295581Z","shell.execute_reply":"2026-01-10T09:18:03.505484Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport esm\nimport os\nfrom tqdm import tqdm\n\n# Use dtype instead of deprecated torch_dtype\n# Updated to utilize both T4 GPUs on Kaggle\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ntorch.set_grad_enabled(False)\n\n# Load model weights\nmodel, alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n\n# Set model to half precision (float16) and move to device\nmodel = model.to(device=device, dtype=torch.float16)\n\n# Wrap model for Multi-GPU utilization (T4 x2)\nif torch.cuda.device_count() > 1:\n    print(f\"Using {torch.cuda.device_count()} GPUs for extraction.\")\n    model = torch.nn.DataParallel(model)\n\nmodel.eval()\nbatch_converter = alphabet.get_batch_converter()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:18:50.398757Z","iopub.execute_input":"2026-01-10T09:18:50.399073Z","iopub.status.idle":"2026-01-10T09:19:10.092593Z","shell.execute_reply.started":"2026-01-10T09:18:50.399046Z","shell.execute_reply":"2026-01-10T09:19:10.091884Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_esm2_embeddings(data, batch_size=2, max_len=1024):\n    \"\"\"\n    Extracts embeddings using DataParallel across multiple GPUs.\n    Truncates sequences to prevent quadratic memory OOM on individual T4 cards.\n    \"\"\"\n    embeddings = {}\n\n    for i in tqdm(range(0, len(data), batch_size), desc=\"Extracting ESM2 embeddings\"):\n        # Slice batch and apply truncation logic for hardware stability\n        raw_batch = data[i:i+batch_size]\n        truncated_batch = [(pid, seq[:max_len]) for pid, seq in raw_batch]\n\n        labels, strs, toks = batch_converter(truncated_batch)\n        toks = toks.to(device)\n\n        with torch.no_grad():\n            # DataParallel automatically splits the batch across GPU 0 and GPU 1\n            out = model(toks, repr_layers=[33], return_contacts=False)\n\n        # Extract representations from the last layer (33)\n        reps = out[\"representations\"][33]\n\n        for j, (protein_id, original_seq) in enumerate(raw_batch):\n            # Account for truncation in indexing\n            seq_len = min(len(original_seq), max_len)\n            \n            # Slice representation (1:seq_len+1) to remove BOS/EOS tokens\n            rep = reps[j, 1:seq_len+1]\n            \n            # Mean pooling and conversion to float16 to save RAM\n            emb = rep.mean(0).cpu().numpy().astype(np.float16)\n            embeddings[protein_id] = emb\n\n        # Regular cache clearing to prevent fragmentation on dual T4s\n        if i % 50 == 0:\n            torch.cuda.empty_cache()\n\n    return embeddings","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:19:30.408842Z","iopub.execute_input":"2026-01-10T09:19:30.409468Z","iopub.status.idle":"2026-01-10T09:19:30.416529Z","shell.execute_reply.started":"2026-01-10T09:19:30.409437Z","shell.execute_reply":"2026-01-10T09:19:30.415888Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_embeddings = extract_esm2_embeddings(train_data, batch_size=2)\nX_train = np.vstack([train_embeddings[pid] for pid in df_train[\"protein_id\"]])\nnp.save(\"X_train_esm2.npy\", X_train)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T09:20:55.441882Z","iopub.execute_input":"2026-01-10T09:20:55.442206Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_embeddings = extract_esm2_embeddings(test_data, batch_size=2)\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,"execution":{"iopub.status.busy":"2026-01-02T05:34:07.694854Z","iopub.execute_input":"2026-01-02T05:34:07.695154Z","iopub.status.idle":"2026-01-02T05:34:25.150067Z","shell.execute_reply.started":"2026-01-02T05:34:07.695128Z","shell.execute_reply":"2026-01-02T05:34:25.148902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}