{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":116062,"databundleVersionId":14875579},{"sourceType":"datasetVersion","sourceId":13582328,"datasetId":8629066,"databundleVersionId":14311941}],"dockerImageVersionId":31153,"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-03-02T10:43:42.426335Z","iopub.execute_input":"2026-03-02T10:43:42.426875Z","iopub.status.idle":"2026-03-02T10:43:42.747001Z","shell.execute_reply.started":"2026-03-02T10:43:42.426847Z","shell.execute_reply":"2026-03-02T10:43:42.746138Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:43:47.751454Z","iopub.execute_input":"2026-03-02T10:43:47.752341Z","iopub.status.idle":"2026-03-02T10:43:47.755875Z","shell.execute_reply.started":"2026-03-02T10:43:47.752312Z","shell.execute_reply":"2026-03-02T10:43:47.755155Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_tax = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\", sep=\"\\t\", header=None,\n names=[\"protein_id\", \"taxon_id\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:43:50.911333Z","iopub.execute_input":"2026-03-02T10:43:50.911862Z","iopub.status.idle":"2026-03-02T10:43:50.961892Z","shell.execute_reply.started":"2026-03-02T10:43:50.91183Z","shell.execute_reply":"2026-03-02T10:43:50.961018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_tax","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:43:53.741602Z","iopub.execute_input":"2026-03-02T10:43:53.742113Z","iopub.status.idle":"2026-03-02T10:43:53.752898Z","shell.execute_reply.started":"2026-03-02T10:43:53.742082Z","shell.execute_reply":"2026-03-02T10:43:53.75215Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(df_tax.nunique())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:43:57.276439Z","iopub.execute_input":"2026-03-02T10:43:57.276758Z","iopub.status.idle":"2026-03-02T10:43:57.297135Z","shell.execute_reply.started":"2026-03-02T10:43:57.276734Z","shell.execute_reply":"2026-03-02T10:43:57.296552Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train_terms = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv',\n sep='\\t')\ndf_train_terms.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:43:59.811386Z","iopub.execute_input":"2026-03-02T10:43:59.811674Z","iopub.status.idle":"2026-03-02T10:44:00.054072Z","shell.execute_reply.started":"2026-03-02T10:43:59.811653Z","shell.execute_reply":"2026-03-02T10:44:00.053449Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train_terms.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:05.572093Z","iopub.execute_input":"2026-03-02T10:44:05.572369Z","iopub.status.idle":"2026-03-02T10:44:05.577753Z","shell.execute_reply.started":"2026-03-02T10:44:05.572345Z","shell.execute_reply":"2026-03-02T10:44:05.576983Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train_terms.nunique()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:06.276908Z","iopub.execute_input":"2026-03-02T10:44:06.277179Z","iopub.status.idle":"2026-03-02T10:44:06.396113Z","shell.execute_reply.started":"2026-03-02T10:44:06.277156Z","shell.execute_reply":"2026-03-02T10:44:06.39533Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_ia = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/IA.tsv', sep='\\t', header=None, names=['term','ia'])\n\ndf_ia","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:07.138378Z","iopub.execute_input":"2026-03-02T10:44:07.139033Z","iopub.status.idle":"2026-03-02T10:44:07.178949Z","shell.execute_reply.started":"2026-03-02T10:44:07.139005Z","shell.execute_reply":"2026-03-02T10:44:07.178217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install Bio","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:12.252044Z","iopub.execute_input":"2026-03-02T10:44:12.252635Z","iopub.status.idle":"2026-03-02T10:44:13.957216Z","shell.execute_reply.started":"2026-03-02T10:44:12.2526Z","shell.execute_reply":"2026-03-02T10:44:13.956308Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nfrom Bio import SeqIO\n\nrecords = []\n\nfor record in SeqIO.parse(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\", \"fasta\"):\n header = record.description # full header line, everything after '>'\n \n # Split header on \"|\" to extract fields\n parts = header.split(\"|\")\n \n # Usually in UniProt FASTA:\n # parts[0] = prefix like \"sp\" or \"tr\"\n # parts[1] = accession like \"A0A0C5B5G6\"\n # parts[2] = entry name like \"MOTSC_HUMAN\"\n # Everything after is description (if any)\n \n prefix = parts[0] if len(parts) > 0 else \"\"\n accession = parts[1] if len(parts) > 1 else \"\"\n entry_name = parts[2] if len(parts) > 2 else \"\"\n \n seq = str(record.seq)\n \n records.append({ \n \"db_prefix\": prefix, \n \"accession\": accession, \n \"entry_name\": entry_name, \n \"sequence\": seq\n })\n\ndf_fasta0 = pd.DataFrame(records)\n\ndf_fasta0.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:19.532111Z","iopub.execute_input":"2026-03-02T10:44:19.532849Z","iopub.status.idle":"2026-03-02T10:44:20.163122Z","shell.execute_reply.started":"2026-03-02T10:44:19.532809Z","shell.execute_reply":"2026-03-02T10:44:20.16238Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_fasta0['entry_name']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:20.356412Z","iopub.execute_input":"2026-03-02T10:44:20.356706Z","iopub.status.idle":"2026-03-02T10:44:20.363967Z","shell.execute_reply.started":"2026-03-02T10:44:20.356684Z","shell.execute_reply":"2026-03-02T10:44:20.363139Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_fasta0.loc[6, 'entry_name']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:24.326658Z","iopub.execute_input":"2026-03-02T10:44:24.326958Z","iopub.status.idle":"2026-03-02T10:44:24.332709Z","shell.execute_reply.started":"2026-03-02T10:44:24.326937Z","shell.execute_reply":"2026-03-02T10:44:24.331948Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\nimport pandas as pd\nfrom Bio import SeqIO\n\nrecords = []\n\nfor record in SeqIO.parse(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\", \"fasta\"):\n header = record.description.strip()\n sequence = str(record.seq)\n \n # Pattern to match the UniProt FASTA metadata\n pattern = r\"^(?Psp|tr)\\|(?P[^|]+)\\|(?P[^ ]+)\\s(?P.*?)(?:\\sOS=(?P.*?))?(?:\\sOX=(?P\\d+))?(?:\\sGN=(?P[^ ]+))?(?:\\sPE=(?P\\d+))?(?:\\sSV=(?P\\d+))?$\"\n match = re.match(pattern, header)\n \n if match:\n info = match.groupdict()\n else:\n # Fallback if pattern fails\n info = {\n \"prefix\": \"\",\n \"accession\": \"\",\n \"entry\": \"\",\n \"description\": header,\n \"organism\": \"\",\n \"taxon\": \"\",\n \"gene\": \"\",\n \"pe\": \"\",\n \"sv\": \"\"\n }\n \n info[\"sequence\"] = sequence\n records.append(info)\n\ndf_fasta = pd.DataFrame(records)\ndf_fasta.head(10)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:27.616708Z","iopub.execute_input":"2026-03-02T10:44:27.616981Z","iopub.status.idle":"2026-03-02T10:44:29.628121Z","shell.execute_reply.started":"2026-03-02T10:44:27.616959Z","shell.execute_reply":"2026-03-02T10:44:29.62735Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_fasta.isnull().sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:29.629669Z","iopub.execute_input":"2026-03-02T10:44:29.630016Z","iopub.status.idle":"2026-03-02T10:44:29.697252Z","shell.execute_reply.started":"2026-03-02T10:44:29.629994Z","shell.execute_reply":"2026-03-02T10:44:29.696484Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"missing_gene = df_fasta[df_fasta['gene'].isnull()]\nmissing_gene","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:36.79149Z","iopub.execute_input":"2026-03-02T10:44:36.792225Z","iopub.status.idle":"2026-03-02T10:44:36.812807Z","shell.execute_reply.started":"2026-03-02T10:44:36.792183Z","shell.execute_reply":"2026-03-02T10:44:36.812044Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"missing_taxon = df_fasta[df_fasta['taxon'].isnull()]\nmissing_taxon","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:43.102371Z","iopub.execute_input":"2026-03-02T10:44:43.103075Z","iopub.status.idle":"2026-03-02T10:44:43.124174Z","shell.execute_reply.started":"2026-03-02T10:44:43.103046Z","shell.execute_reply":"2026-03-02T10:44:43.12355Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"missing_any = df_fasta[df_fasta.isnull().any(axis=1)]\nmissing_any","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:49.894007Z","iopub.execute_input":"2026-03-02T10:44:49.894287Z","iopub.status.idle":"2026-03-02T10:44:49.971423Z","shell.execute_reply.started":"2026-03-02T10:44:49.894265Z","shell.execute_reply":"2026-03-02T10:44:49.970613Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. Select rows with missing taxon or gene\nmissing_mask = df_fasta['taxon'].isna() \ndf_fix = df_fasta[missing_mask].copy()\ndf_fix","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:44:55.651683Z","iopub.execute_input":"2026-03-02T10:44:55.652146Z","iopub.status.idle":"2026-03-02T10:44:55.675266Z","shell.execute_reply.started":"2026-03-02T10:44:55.652113Z","shell.execute_reply":"2026-03-02T10:44:55.674513Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2. Extract OX and GN from organism column\ndf_fix['taxon'] = df_fix['organism'].str.extract(r\"OX=(\\d+)\")\ndf_fix['gene'] = df_fix['organism'].str.extract(r\"GN=([A-Za-z0-9\\-_. ]+)\")\ndf_fix","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:00.71178Z","iopub.execute_input":"2026-03-02T10:45:00.712391Z","iopub.status.idle":"2026-03-02T10:45:00.72808Z","shell.execute_reply.started":"2026-03-02T10:45:00.712366Z","shell.execute_reply":"2026-03-02T10:45:00.727233Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3. Clean organism column\ndf_fix['organism'] = df_fix['organism'].str.replace(r\"OX=\\d+\", \"\", regex=True)\ndf_fix['organism'] = df_fix['organism'].str.replace(r\"GN=[A-Za-z0-9\\-_. ]+\", \"\", regex=True)\ndf_fix['organism'] = df_fix['organism'].str.strip()\ndf_fix","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:06.408375Z","iopub.execute_input":"2026-03-02T10:45:06.409022Z","iopub.status.idle":"2026-03-02T10:45:06.425368Z","shell.execute_reply.started":"2026-03-02T10:45:06.408992Z","shell.execute_reply":"2026-03-02T10:45:06.424407Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 4. Update original df\ndf_fasta.update(df_fix)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:07.633563Z","iopub.execute_input":"2026-03-02T10:45:07.634236Z","iopub.status.idle":"2026-03-02T10:45:07.740579Z","shell.execute_reply.started":"2026-03-02T10:45:07.634213Z","shell.execute_reply":"2026-03-02T10:45:07.739738Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_fasta.isnull().sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:13.228399Z","iopub.execute_input":"2026-03-02T10:45:13.229148Z","iopub.status.idle":"2026-03-02T10:45:13.299027Z","shell.execute_reply.started":"2026-03-02T10:45:13.229109Z","shell.execute_reply":"2026-03-02T10:45:13.298224Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_fasta[df_fasta['gene'].isnull()]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:15.33341Z","iopub.execute_input":"2026-03-02T10:45:15.333695Z","iopub.status.idle":"2026-03-02T10:45:15.355054Z","shell.execute_reply.started":"2026-03-02T10:45:15.333674Z","shell.execute_reply":"2026-03-02T10:45:15.354375Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install obonet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:18.753348Z","iopub.execute_input":"2026-03-02T10:45:18.753637Z","iopub.status.idle":"2026-03-02T10:45:20.375737Z","shell.execute_reply.started":"2026-03-02T10:45:18.753611Z","shell.execute_reply":"2026-03-02T10:45:20.374815Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import obonet\ngraph = obonet.read_obo('/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo')\nuniverse = set(graph.nodes())\nuntrained_universe = set(df_train_terms['term']) - universe\nprint(\"The untrained universe is: \", len(untrained_universe))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:20.377544Z","iopub.execute_input":"2026-03-02T10:45:20.377785Z","iopub.status.idle":"2026-03-02T10:45:25.602429Z","shell.execute_reply.started":"2026-03-02T10:45:20.377762Z","shell.execute_reply":"2026-03-02T10:45:25.601768Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_test = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\",sep='\\t')\ndf_test","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:25.60344Z","iopub.execute_input":"2026-03-02T10:45:25.603971Z","iopub.status.idle":"2026-03-02T10:45:25.624772Z","shell.execute_reply.started":"2026-03-02T10:45:25.603949Z","shell.execute_reply":"2026-03-02T10:45:25.62383Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\nimport pandas as pd\nfrom Bio import SeqIO\n\ndef parse_uniprot_header(header):\n \"\"\"\n Fully robust parser that handles both UniProt-style and non-standard FASTA headers.\n \"\"\"\n\n # ----------------------------------------------------\n # CASE 1: UniProt-style header: sp|Q9SZZ8|BCH1_ARATH ...\n # ----------------------------------------------------\n if \"|\" in header and header.count(\"|\") >= 2:\n main, *rest = header.split(\" \", 1)\n parts = main.split(\"|\")\n\n if len(parts) == 3:\n prefix, accession, entry = parts\n else:\n # fallback if malformed\n prefix, accession, entry = None, None, None\n\n remaining = rest[0] if rest else \"\"\n\n else:\n # ----------------------------------------------------\n # CASE 2: Non-UniProt header\n # we treat entire header as description\n # ----------------------------------------------------\n return {\n \"prefix\": None,\n \"accession\": None,\n \"entry\": None,\n \"description\": header,\n \"organism\": None,\n \"taxon\": None,\n \"gene\": None,\n \"pe\": None,\n \"sv\": None\n }\n\n # --------------------------\n # Extract metadata tags\n # --------------------------\n\n def extract_tag(pattern):\n m = re.search(pattern, header)\n return m.group(1).strip() if m else None\n\n organism = extract_tag(r\"OS=(.+?)(?=\\s[A-Z]{2}=|$)\")\n taxon = extract_tag(r\"OX=(\\d+)\")\n gene = extract_tag(r\"GN=(.+?)(?=\\s[A-Z]{2}=|$)\")\n pe = extract_tag(r\"PE=(\\d+)\")\n sv = extract_tag(r\"SV=(\\d+)\")\n\n # --------------------------\n # Clean description\n # --------------------------\n description = re.split(r\"\\s[A-Z]{2}=\", remaining)[0].strip()\n\n return {\n \"prefix\": prefix,\n \"accession\": accession,\n \"entry\": entry,\n \"description\": description,\n \"organism\": organism,\n \"taxon\": taxon,\n \"gene\": gene,\n \"pe\": pe,\n \"sv\": sv\n }\n\n\n# --------------------------\n# MAIN LOOP\n# --------------------------\nrecords = []\nfasta_file = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\nfor record in SeqIO.parse(fasta_file, \"fasta\"):\n header = record.description.strip()\n seq = str(record.seq)\n\n info = parse_uniprot_header(header)\n info[\"sequence\"] = seq\n\n records.append(info)\n\ndf_test_fasta = pd.DataFrame(records)\ndf_test_fasta.head(20)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:35.549176Z","iopub.execute_input":"2026-03-02T10:45:35.549838Z","iopub.status.idle":"2026-03-02T10:45:37.181851Z","shell.execute_reply.started":"2026-03-02T10:45:35.549806Z","shell.execute_reply":"2026-03-02T10:45:37.181083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio import SeqIO\nimport pandas as pd\n\ndef parse_test_header(header):\n \"\"\"\n Parse headers of format: '>ACCESSION TAXON'\n Example: '>A0PK11 9606'\n \"\"\"\n header = header.lstrip(\">\") # remove leading >\n parts = header.split()\n\n accession = parts[0] if len(parts) > 0 else None\n taxon = parts[1] if len(parts) > 1 else None\n\n return {\n \"accession\": accession,\n \"taxon\": taxon\n }\n\ndef load_test_fasta(path):\n rows = []\n\n for record in SeqIO.parse(path, \"fasta\"):\n header = record.description\n parsed = parse_test_header(header)\n\n row = {\n \"accession\": parsed[\"accession\"],\n \"taxon\": parsed[\"taxon\"],\n \"sequence\": str(record.seq)\n }\n\n rows.append(row)\n\n return pd.DataFrame(rows)\n\n# --------------------------\n# šŸ”§ USAGE\n# --------------------------\ndf_test_fasta = load_test_fasta(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\nprint(df_test.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:37.183193Z","iopub.execute_input":"2026-03-02T10:45:37.183418Z","iopub.status.idle":"2026-03-02T10:45:38.400155Z","shell.execute_reply.started":"2026-03-02T10:45:37.183399Z","shell.execute_reply":"2026-03-02T10:45:38.399412Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nfrom Bio import SeqIO\n\n# Load training GO annotations\ndf_train_terms = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv', sep='\\t')\n\n# Load test accessions from testsuperset.fasta\ntest_accessions = []\nfor record in SeqIO.parse('/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta', 'fasta'):\n # header format is like \"A0A0C5B5G6 9606\"\n accession = record.description.split()[0]\n test_accessions.append(accession)\n\ntest_accessions = set(test_accessions)\n\n# Which of these appear in training\ntrain_terms_accessions = set(df_train_terms['EntryID'])\noverlap_proteins = test_accessions.intersection(train_terms_accessions)\n\nprint(\"Number of test proteins that also appear in training terms:\", len(overlap_proteins))\nprint(\"Sample overlap proteins:\", list(overlap_proteins)[:10])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:38.401304Z","iopub.execute_input":"2026-03-02T10:45:38.4017Z","iopub.status.idle":"2026-03-02T10:45:39.581794Z","shell.execute_reply.started":"2026-03-02T10:45:38.401669Z","shell.execute_reply":"2026-03-02T10:45:39.581112Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\n# load training terms\ndf_train_terms = pd.read_csv(\n \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n sep=\"\\t\",\n names=[\"EntryID\", \"term\", \"aspect\"]\n)\n\n# load test superset taxon list to get test IDs\ndf_test_taxon = pd.read_csv(\n \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\",\n sep=\"\\t\",\n names=[\"protein_id\", \"taxon_id\"]\n)\n\n# get unique protein sets\ntrain_proteins = set(df_train_terms[\"EntryID\"].unique())\ntest_proteins = set(df_test_taxon[\"protein_id\"].unique())\n\n# compute overlap\ncommon_proteins = test_proteins.intersection(train_proteins)\n\nprint(\"Number of test proteins:\", len(test_proteins))\nprint(\"Number of test proteins also in training terms:\", len(common_proteins))\n\n# show some examples\nprint(\"Some overlapping proteins:\", list(common_proteins)[:10])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:44.643337Z","iopub.execute_input":"2026-03-02T10:45:44.643867Z","iopub.status.idle":"2026-03-02T10:45:44.918999Z","shell.execute_reply.started":"2026-03-02T10:45:44.643839Z","shell.execute_reply":"2026-03-02T10:45:44.918165Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Visualize a part of the GO graph (go-basic.obo)\n# pip install obonet networkx matplotlib pygraphviz # pygraphviz optional (better layouts)\n# or pip install pydot and graphviz system package (if pygraphviz fails)\n\nimport obonet\nimport networkx as nx\nimport matplotlib.pyplot as plt\nfrom textwrap import shorten\n\ndef load_go(path=\"go-basic.obo\"):\n \"\"\"Load GO OBO into a NetworkX DiGraph using obonet.\"\"\"\n G = obonet.read_obo(path)\n # obonet returns a DiGraph where nodes have attributes like 'name', 'namespace'\n return G\n\ndef get_subgraph_by_root(G, root_go, depth=20):\n \"\"\"\n Return a subgraph containing nodes up to `depth` edges below root_go\n (i.e., children, grandchildren, ... up to `depth`).\n This works regardless of the edge orientation in the OBO file by using\n the reversed graph for descendant traversal.\n \"\"\"\n if root_go not in G:\n raise KeyError(f\"{root_go} not in GO graph\")\n\n # reverse graph so that descendants of root_go correspond to more specific terms\n Gr = G.reverse(copy=True)\n\n # single_source_shortest_path_length gives distance (in edges); cutoff=depth\n nodes_within = nx.single_source_shortest_path_length(Gr, root_go, cutoff=depth)\n nodes = set(nodes_within.keys()) | {root_go}\n sub = G.subgraph(nodes).copy()\n return sub\n\ndef make_labels(subG, show_id=True, maxlen=30):\n \"\"\"Return dict of node->label to show on plot (shortened name + optional GO id).\"\"\"\n labels = {}\n for n, data in subG.nodes(data=True):\n name = data.get(\"name\", \"\")\n if show_id:\n label = f\"{n}\\\\n{shorten(name, width=maxlen)}\"\n else:\n label = shorten(name, width=maxlen)\n labels[n] = label\n return labels\n\ndef draw_go_subgraph(subG, labels=None, figsize=(12,10)):\n \"\"\"Plot a NetworkX DiGraph of GO terms.\"\"\"\n plt.figure(figsize=figsize)\n\n # try graphviz layout for DAG-like layout\n try:\n pos = nx.nx_agraph.graphviz_layout(subG, prog=\"dot\") # requires pygraphviz / pydot + graphviz\n except Exception:\n # fallback to spring layout (less tidy for hierarchies)\n pos = nx.spring_layout(subG, k=0.6, iterations=200)\n\n # draw nodes and edges\n nx.draw_networkx_nodes(subG, pos, node_size=300, node_shape=\"o\")\n nx.draw_networkx_edges(subG, pos, arrows=True, arrowsize=12, width=1.0)\n if labels is None:\n labels = make_labels(subG)\n nx.draw_networkx_labels(subG, pos, labels=labels, font_size=8)\n\n plt.axis(\"off\")\n plt.title(\"GO subgraph visualization\")\n plt.tight_layout()\n plt.show()\n\nif __name__ == \"__main__\":\n # path to your OBO file\n obo_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n\n G = load_go(obo_path)\n\n # Example roots (as provided in the dataset)\n subontology_roots = {\n 'BPO': 'GO:0008150', # Biological Process\n 'CCO': 'GO:0005575', # Cellular Component\n 'MFO': 'GO:0003674' # Molecular Function\n }\n\n # Choose which root/term to visualize and depth\n root = subontology_roots['MFO'] # try 'BPO' or 'CCO' or any GO:XXXXXXX id\n depth = 2 # increase to see more levels (beware of huge graphs)\n\n subG = get_subgraph_by_root(G, root, depth=depth)\n labels = make_labels(subG, show_id=True, maxlen=40)\n draw_go_subgraph(subG, labels=labels)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:45:45.013086Z","iopub.execute_input":"2026-03-02T10:45:45.013535Z","iopub.status.idle":"2026-03-02T10:45:54.33785Z","shell.execute_reply.started":"2026-03-02T10:45:45.013505Z","shell.execute_reply":"2026-03-02T10:45:54.337131Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":" !pip install transformers torch scikit-learn pandas tqdm tensorflow","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:46:23.569342Z","iopub.execute_input":"2026-03-02T10:46:23.569707Z","iopub.status.idle":"2026-03-02T10:46:26.328232Z","shell.execute_reply.started":"2026-03-02T10:46:23.569677Z","shell.execute_reply":"2026-03-02T10:46:26.327398Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install \"numpy<2.1.0\" \"scipy<1.14.0\" \"pandas==2.2.2\" \"scikit-learn<1.7\"\n!pip install google-auth==2.38.0 notebook==6.5.7 requests==2.32.3 tornado==6.4.2","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install \"google-cloud-bigquery>=3.31.0\" \"google-cloud-bigquery-storage>=2.30.0\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell A: reinstall numpy first, then scikit-learn (and scipy if used)\n\n!pip install -q --force-reinstall \"scipy\" \"scikit-learn\"\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --force-reinstall \"numpy<2.0.0\" \"scipy<1.14.0\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom transformers import AutoTokenizer, EsmModel\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.multioutput import MultiOutputClassifier\nfrom sklearn.linear_model import LogisticRegression\nfrom tqdm import tqdm\nimport pickle\n\n\nprint(\"āœ“ All packages imported successfully\")\nprint(f\"PyTorch version: {torch.__version__}\")\nprint(f\"CUDA available: {torch.cuda.is_available()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:46:42.103836Z","iopub.execute_input":"2026-03-02T10:46:42.104148Z","iopub.status.idle":"2026-03-02T10:47:12.457433Z","shell.execute_reply.started":"2026-03-02T10:46:42.10412Z","shell.execute_reply":"2026-03-02T10:47:12.456545Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install torch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T10:56:08.754045Z","iopub.execute_input":"2026-03-02T10:56:08.755305Z","iopub.status.idle":"2026-03-02T10:56:10.473063Z","shell.execute_reply.started":"2026-03-02T10:56:08.755262Z","shell.execute_reply":"2026-03-02T10:56:10.47202Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport pandas as pd\nfrom tqdm import tqdm\nfrom transformers import AutoTokenizer, AutoModel\nfrom Bio import SeqIO\nimport time\nimport math\n\n# ============ Settings ============\n# (kept same as your original script)\nMODEL_NAME = \"Rostlab/prot_bert_bfd\" # ProtBERT\n\nTRAIN_FASTA = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\nTRAIN_TERMS = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\nTRAIN_TAXON = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\"\nIA_WEIGHTS = \"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\"\n\nOUT_EMB_DIR = \"/kaggle/working/protbert_embeddings\"\nos.makedirs(OUT_EMB_DIR, exist_ok=True)\n\n# checkpointing frequency: number of sequences processed before saving a partial checkpoint\nCHECKPOINT_FREQ = 5000\n\n# starting batch size (will drop automatically on OOM)\nBATCH_SIZE = 16\n\n# GPU setup\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Using device:\", device)\n\n# ============ Load Data ============\n\n# Parse sequences\nseq_records = []\nfor record in SeqIO.parse(TRAIN_FASTA, \"fasta\"):\n seq_records.append({\n \"accession\": record.id,\n \"sequence\": str(record.seq)\n })\ndf_seqs = pd.DataFrame(seq_records)\nprint(\"Total sequences:\", len(df_seqs))\n\n# Load train_terms (annotations)\ndf_terms = pd.read_csv(TRAIN_TERMS, sep=\"\\t\", names=[\"accession\",\"go_term\",\"aspect\"])\nprint(\"Train terms shape:\", df_terms.shape)\n\n# Load taxonomy\ndf_taxon = pd.read_csv(TRAIN_TAXON, sep=\"\\t\", names=[\"accession\",\"taxon_id\"])\nprint(\"Taxonomy shape:\", df_taxon.shape)\n\n# Load IA weights\ndf_ia = pd.read_csv(IA_WEIGHTS, sep=\"\\t\", names=[\"go_term\",\"ia\"])\nprint(\"IA weights shape:\", df_ia.shape)\n\n# Merge all into a master training table\ndf_master = df_seqs.merge(df_taxon, on=\"accession\", how=\"left\")\nprint(\"Master DF shape:\", df_master.shape)\n\n# ============ ProtBERT Model ============\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, do_lower_case=False)\n# set a safe model_max_length to avoid the truncation warning (ProtBERT accepts long inputs; 1024 is safe)\nif getattr(tokenizer, \"model_max_length\", None) is None or tokenizer.model_max_length > 10000:\n tokenizer.model_max_length = 1024\n\nmodel = AutoModel.from_pretrained(MODEL_NAME)\n\n# Use both GPUs if available (keeps your original design)\nif torch.cuda.device_count() > 1:\n print(f\"Using {torch.cuda.device_count()} GPUs!\")\n model = torch.nn.DataParallel(model)\n\nmodel.to(device)\nmodel.eval()\n\n# ============ Embedding Helper (with checkpointing + OOM handling) ============\n\ndef embed_batch(sequences, batch_size=16, device=device):\n \"\"\"\n Embed a list of protein sequences with ProtBERT.\n This function uses mixed precision and handles OOM by retrying with half batch size.\n Returns: tensor of shape (len(sequences), hidden_size) on CPU\n \"\"\"\n all_outputs = []\n b = 0\n current_batch_size = batch_size\n\n while b < len(sequences):\n end = min(b + current_batch_size, len(sequences))\n batch = sequences[b:end]\n # ProtBERT expects sequences space-separated like \"M A T K V ...\"\n batch_input = [\" \".join(list(seq)) for seq in batch]\n\n # Tokenize on CPU to keep GPU memory low\n inputs = tokenizer(batch_input, return_tensors=\"pt\", padding=True, truncation=True, max_length=tokenizer.model_max_length)\n # Move inputs to device\n inputs = {k: v.to(device) for k, v in inputs.items()}\n\n try:\n with torch.no_grad():\n # mixed precision to reduce memory\n with torch.cuda.amp.autocast(enabled=(device != \"cpu\")):\n outputs = model(**inputs)\n\n # take CLS embedding (first token)\n cls_emb = outputs.last_hidden_state[:, 0, :].detach().cpu()\n all_outputs.append(cls_emb)\n\n # free GPU memory asap\n del outputs\n del inputs\n torch.cuda.empty_cache()\n\n # advance\n b = end\n\n # if we succeeded and the batch size is smaller than requested, try to grow it slowly (optional)\n # but keep this conservative: don't change for stability\n except RuntimeError as e:\n # catch CUDA OOM only\n if \"out of memory\" in str(e).lower():\n print(f\"CUDA OOM with batch size {current_batch_size} at index {b}. Reducing batch size and retrying...\")\n # free cache and reduce batch\n del inputs\n torch.cuda.empty_cache()\n if current_batch_size <= 1:\n # cannot reduce further; re-raise\n raise\n current_batch_size = max(1, current_batch_size // 2)\n print(\"New batch size:\", current_batch_size)\n # small sleep to let CUDA driver settle\n time.sleep(2)\n continue\n else:\n # unexpected error -> reraise\n raise\n\n if len(all_outputs) == 0:\n return torch.empty((0, model.module.config.hidden_size if hasattr(model, \"module\") else model.config.hidden_size))\n return torch.cat(all_outputs, dim=0)\n\n# ============ Compute & Save Embeddings with periodic checkpoints ============\n\nprint(\"Generating ProtBERT embeddings with checkpointing...\")\n\nall_seqs = df_master[\"sequence\"].tolist()\nall_accessions = df_master[\"accession\"].tolist()\nn = len(all_seqs)\n\n# We'll process in chunks of CHECKPOINT_FREQ and save after each chunk\nembeds_parts = [] # list of tensors on CPU\nprocessed_ids = []\n\n# chunking loop\nfor chunk_start in range(0, n, CHECKPOINT_FREQ):\n chunk_end = min(chunk_start + CHECKPOINT_FREQ, n)\n chunk_seqs = all_seqs[chunk_start:chunk_end]\n chunk_ids = all_accessions[chunk_start:chunk_end]\n\n print(f\"Embedding chunk {chunk_start} -> {chunk_end} (size {len(chunk_seqs)})\")\n\n # embed chunk (this returns cpu tensor)\n emb_chunk = embed_batch(chunk_seqs, batch_size=BATCH_SIZE, device=device)\n # sanity check: emb_chunk.shape[0] should equal len(chunk_seqs)\n if emb_chunk.shape[0] != len(chunk_seqs):\n # this should not happen; but if it does, pad with zeros for missing items\n print(\"Warning: embeddings count mismatch; padding to match sequences\")\n missing = len(chunk_seqs) - emb_chunk.shape[0]\n pad = torch.zeros((missing, emb_chunk.shape[1]))\n emb_chunk = torch.cat([emb_chunk, pad], dim=0)\n\n embeds_parts.append(emb_chunk)\n processed_ids.extend(chunk_ids)\n\n # checkpoint save\n ckpt_path = os.path.join(OUT_EMB_DIR, f\"train_protbert_embeddings_ckpt_{chunk_end}.pt\")\n torch.save({\n \"accessions\": processed_ids,\n \"embeddings\": torch.cat(embeds_parts, dim=0)\n }, ckpt_path)\n print(f\"Saved checkpoint to {ckpt_path}\")\n\n # free memory references\n del emb_chunk\n torch.cuda.empty_cache()\n\n# final concatenation & save\nfinal_embeddings = torch.cat(embeds_parts, dim=0)\nemb_file = os.path.join(OUT_EMB_DIR, \"train_protbert_embeddings.pt\")\ntorch.save({\n \"accessions\": processed_ids,\n \"embeddings\": final_embeddings\n}, emb_file)\n\nprint(\"Saved ProtBERT embeddings to:\", emb_file)\n\n# ============ Optional: Save as CSV for metadata ============\n\ndf_meta = df_master[[\"accession\",\"taxon_id\"]]\nmeta_file = os.path.join(OUT_EMB_DIR, \"train_metadata.csv\")\ndf_meta.to_csv(meta_file, index=False)\n\nprint(\"Saved metadata for embeddings to:\", meta_file)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T11:29:45.958577Z","iopub.execute_input":"2026-03-02T11:29:45.959082Z","iopub.status.idle":"2026-03-02T12:30:46.008512Z","shell.execute_reply.started":"2026-03-02T11:29:45.959052Z","shell.execute_reply":"2026-03-02T12:30:46.007658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # cell: deepgozero_el_embeddings_pipeline.py\n# # RUN THIS IN A KAGGLE NOTEBOOK CELL (one cell).\n# # It will:\n# # - clone deepgozero + el-embeddings repos\n# # - download DeepGO data tarball (if available)\n# # - try to run the EL normalization (Normalizer.groovy) and elembedding.py in chunks\n# # - save checkpointed class embeddings every chunk\n# # - if EL pipeline fails due to missing system deps (groovy/jcel), fall back to node2vec on GO OBO\n# #\n# # NOTE: Running the full EL embedding training is computationally intensive and requires JDK/groovy.\n# # The fallback node2vec provides a workable embedding matrix quickly.\n\n# import os\n# import sys\n# import subprocess\n# import logging\n# import time\n# import shutil\n# from pathlib import Path\n# import pickle\n\n# import numpy as np\n# import obonet\n# import networkx as nx\n\n# # ---------- Config ----------\n# ROOT = Path(\"/kaggle/working/deepgozero_el_embeddings\")\n# ROOT.mkdir(parents=True, exist_ok=True)\n\n# LOG_FILE = ROOT / \"el_embedding_run.log\"\n\n# # Repos and data URLs (public)\n# DEEPGOZERO_GIT = \"https://github.com/bio-ontology-research-group/deepgozero.git\"\n# EL_EMBED_GIT = \"https://github.com/bio-ontology-research-group/el-embeddings.git\"\n# DEEPGO_DATA_BASE = \"http://deepgo.cbrc.kaust.edu.sa/data/deepgozero/\" # index page\n# DEEPGO_DATA_TAR = \"http://deepgo.cbrc.kaust.edu.sa/data/deepgozero/data.tar.gz\"\n\n# # local places\n# REPO_DIR = ROOT / \"repos\"\n# REPO_DIR.mkdir(exist_ok=True)\n# DEEPGOZERO_DIR = REPO_DIR / \"deepgozero\"\n# EL_EMBED_DIR = REPO_DIR / \"el-embeddings\"\n# DATA_DIR = ROOT / \"data\"\n# DATA_DIR.mkdir(exist_ok=True)\n# CHECKPOINT_DIR = ROOT / \"checkpoints\"\n# CHECKPOINT_DIR.mkdir(exist_ok=True)\n\n# # EL training params (adjust for your GPU memory)\n# EL_TOTAL_EPOCHS = 200\n# EL_CHUNK_EPOCHS = 20 # run elembedding.py for 20 epochs at a time, then checkpoint\n# EL_BATCH_SIZE = 256\n# EL_EMBED_SIZE = 200\n# EL_DEVICE = \"cuda\" # we set CUDA_VISIBLE_DEVICES env to \"0,1\" below to attempt both GPUs\n\n# # node2vec fallback params\n# NODE2VEC_DIM = 200\n# NODE2VEC_WALKS = 80\n# NODE2VEC_WINDOW = 10\n# NODE2VEC_EPOCHS = 5\n# NODE2VEC_WORKERS = 2\n\n# # set CUDA devices to both GPUs (so subprocesses see both) -- user said there are 2 GPUs\n# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1\"\n\n# # ---------- Logging ----------\n# logging.basicConfig(\n# level=logging.INFO,\n# format=\"%(asctime)s [%(levelname)s] %(message)s\",\n# handlers=[\n# logging.FileHandler(LOG_FILE),\n# logging.StreamHandler(sys.stdout)\n# ]\n# )\n# log = logging.getLogger(\"deepgozero-el\")\n\n# def run_cmd(cmd, cwd=None, env=None, check=True):\n# \"\"\"Helper to run subprocess commands, capture output in log.\"\"\"\n# log.info(\"RUN: \" + \" \".join(cmd))\n# try:\n# proc = subprocess.run(cmd, cwd=cwd, env=env or os.environ, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, check=check)\n# log.info(proc.stdout[:2000]) # log first chunk of output to keep log compact\n# return True\n# except subprocess.CalledProcessError as e:\n# log.error(\"Command failed with exit code %s\", e.returncode)\n# log.error(e.stdout[:4000] if e.stdout else \"no stdout\")\n# return False\n\n# # ---------- Clone repos ----------\n# def clone_if_missing(url, dest: Path):\n# if dest.exists():\n# log.info(f\"{dest} already exists, skipping clone.\")\n# return True\n# log.info(f\"Cloning {url} -> {dest}\")\n# return run_cmd([\"git\", \"clone\", \"--depth\", \"1\", url, str(dest)])\n\n# # ---------- Download & extract DeepGO data tarball ----------\n# def download_deepgo_data():\n# tar_path = DATA_DIR / \"data.tar.gz\"\n# if tar_path.exists():\n# log.info(\"DeepGO data tarball already exists.\")\n# return True\n# log.info(\"Downloading DeepGO data tarball (this is large).\")\n# if not run_cmd([\"wget\", \"-q\", \"-O\", str(tar_path), DEEPGO_DATA_TAR]):\n# log.warning(\"wget failed; trying curl\")\n# if not run_cmd([\"curl\", \"-sSL\", \"-o\", str(tar_path), DEEPGO_DATA_TAR]):\n# log.error(\"Could not download DeepGO data tarball automatically.\")\n# return False\n# # extract\n# log.info(\"Extracting tarball (may take a while).\")\n# if run_cmd([\"tar\", \"xzf\", str(tar_path), \"-C\", str(DATA_DIR)]):\n# log.info(\"Data extracted under %s\", str(DATA_DIR))\n# return True\n# else:\n# log.error(\"Failed to extract data tarball.\")\n# return False\n\n# # ---------- Try the EL-Embeddings pipeline (DeepGOZero style) ----------\n# def try_el_pipeline(go_input_owl=None):\n# \"\"\"\n# Steps:\n# - Ensure el-embeddings repo exists\n# - Use its Normalizer.groovy (or included jar) to normalize the ontology (OWL input)\n# - Run elembedding.py in epoch chunks to produce checkpoints\n# Returns True if completed successfully and final class embeddings saved.\n# \"\"\"\n# log.info(\"Attempting EL-embeddings pipeline using repo: %s\", EL_EMBED_DIR)\n# if not EL_EMBED_DIR.exists():\n# if not clone_if_missing(EL_EMBED_GIT, EL_EMBED_DIR):\n# log.error(\"Could not clone el-embeddings.\")\n# return False\n\n# # Ensure elembedding.py exists\n# elembedding_script = EL_EMBED_DIR / \"elembedding.py\"\n# if not elembedding_script.exists():\n# log.error(\"elembedding.py not found in el-embeddings repo; aborting EL pipeline.\")\n# return False\n\n# # Prepare normalized OWL using Normalizer.groovy if available\n# normalized_owl = DATA_DIR / \"go_normalized.owl\"\n# normalizer_script = EL_EMBED_DIR / \"Normalizer.groovy\"\n# jar_dir = EL_EMBED_DIR / \"jar\"\n# jcel_jar = jar_dir / \"jcel.jar\"\n\n# if normalizer_script.exists() and shutil.which(\"groovy\"):\n# log.info(\"Running Normalizer.groovy (requires groovy).\")\n# if not run_cmd([\"groovy\", str(normalizer_script), \"-i\", go_input_owl, \"-o\", str(normalized_owl)], cwd=str(EL_EMBED_DIR)):\n# log.error(\"Normalizer failed.\")\n# return False\n# elif jcel_jar.exists():\n# # try to run jcel.jar directly if possible (some repos include jar)\n# log.info(\"jcel.jar detected; trying to run normalization using jcel.jar + groovy fallback not available.\")\n# # Some workflows call java -jar jcel.jar ...; the exact invocation differs ā€” we'll try the common approach in the el-embeddings README:\n# try:\n# # attempt to run Normalizer.groovy via 'groovy' by invoking jdk 'java -jar jcel.jar' won't run the groovy script.\n# # If groovy not installed, we cannot reliably normalize; so report and abort.\n# log.warning(\"Normalizer.groovy exists but 'groovy' not found. EL pipeline requires 'groovy' or external normalization; aborting EL route.\")\n# return False\n# except Exception as e:\n# log.error(\"Normalization attempt failed: %s\", e)\n# return False\n# else:\n# log.error(\"No Normalizer.groovy or jcel.jar found. EL pipeline cannot normalize ontology in this environment.\")\n# return False\n\n# # At this point normalized_owl should exist\n# if not normalized_owl.exists():\n# log.error(\"Normalized OWL not found at %s\", str(normalized_owl))\n# return False\n\n# # Now run elembedding.py in chunks of EL_CHUNK_EPOCHS to produce checkpoints\n# total = EL_TOTAL_EPOCHS\n# chunk = EL_CHUNK_EPOCHS\n# ckpt_counter = 0\n# success = True\n\n# for start in range(0, total, chunk):\n# run_epochs = min(chunk, total - start)\n# ckpt_suffix = f\"epoch_{start+run_epochs}\"\n# ckpt_classes = CHECKPOINT_DIR / f\"classes_{ckpt_suffix}.pkl\"\n# ckpt_rels = CHECKPOINT_DIR / f\"rels_{ckpt_suffix}.pkl\"\n# lossfile = CHECKPOINT_DIR / f\"loss_{ckpt_suffix}.pkl\"\n\n# cmd = [\n# sys.executable,\n# str(elembedding_script),\n# \"--data-file\", str(normalized_owl),\n# \"--out-classes-file\", str(ckpt_classes),\n# \"--out-relations-file\", str(ckpt_rels),\n# \"--epochs\", str(run_epochs),\n# \"--batch-size\", str(EL_BATCH_SIZE),\n# \"--embedding-size\", str(EL_EMBED_SIZE),\n# \"--device\", EL_DEVICE,\n# \"--loss-history-file\", str(lossfile)\n# ]\n# # try to call elembedding.py (it must support these args per repo README)\n# env = os.environ.copy()\n# # ensure both GPUs visible\n# env[\"CUDA_VISIBLE_DEVICES\"] = os.environ.get(\"CUDA_VISIBLE_DEVICES\",\"0,1\")\n# log.info(\"Launching elembedding.py chunk %s\", ckpt_suffix)\n# ok = run_cmd(cmd, cwd=str(EL_EMBED_DIR), env=env)\n# if not ok:\n# log.error(\"elembedding.py failed on chunk %s. Aborting EL pipeline.\", ckpt_suffix)\n# success = False\n# break\n# else:\n# log.info(\"Chunk finished; checkpoint saved: %s\", ckpt_classes)\n# ckpt_counter += 1\n# # small pause\n# time.sleep(1)\n\n# if success:\n# # assume last checkpoint contains final classes\n# final_ckpts = sorted(CHECKPOINT_DIR.glob(\"classes_epoch_*.pkl\")) + sorted(CHECKPOINT_DIR.glob(\"classes_epoch_*.pkl\"))\n# # fallback: find highest epoch file\n# ckpt_files = sorted(CHECKPOINT_DIR.glob(\"classes_*.pkl\"))\n# if ckpt_files:\n# final_ckpt = ckpt_files[-1]\n# final_out = ROOT / \"go_embeddings_deepgozero_classes.pkl\"\n# shutil.copyfile(str(final_ckpt), str(final_out))\n# log.info(\"EL pipeline completed. Final classes embeddings saved to: %s\", final_out)\n# return True\n# else:\n# log.warning(\"EL pipeline ran but no checkpoint classes found.\")\n# return False\n# return False\n\n# # ---------- Fallback: node2vec embeddings on GO OBO ----------\n# def node2vec_fallback(go_obo_path: str, out_file: Path):\n# log.info(\"Building GO graph with obonet and computing node2vec fallback embeddings.\")\n# # parse go\n# G_obo = obonet.read_obo(go_obo_path)\n# G = nx.Graph()\n# for n in G_obo.nodes():\n# G.add_node(n)\n# for u,v,d in G_obo.edges(data=True):\n# # add undirected edge\n# G.add_edge(u, v)\n# log.info(\"GO graph: nodes=%d edges=%d\", G.number_of_nodes(), G.number_of_edges())\n\n# # install node2vec if missing\n# try:\n# from node2vec import Node2Vec\n# except Exception:\n# log.info(\"node2vec python package missing; installing.\")\n# run_cmd([sys.executable, \"-m\", \"pip\", \"install\", \"node2vec\"])\n# from node2vec import Node2Vec\n\n# node2vec = Node2Vec(G, dimensions=NODE2VEC_DIM, walk_length=30, num_walks=NODE2VEC_WALKS, workers=NODE2VEC_WORKERS)\n# model = node2vec.fit(window=NODE2VEC_WINDOW, min_count=1, batch_words=4, epochs=NODE2VEC_EPOCHS)\n\n# embeddings = {n: model.wv[n].astype(np.float32) for n in G.nodes()}\n# with open(out_file, \"wb\") as f:\n# pickle.dump(embeddings, f)\n# log.info(\"Saved node2vec fallback embeddings to %s\", out_file)\n# return True\n\n# # ---------- Main ----------\n# def main():\n# log.info(\"Starting DeepGOZero EL embeddings generation pipeline.\")\n# # 1) clone repos\n# clone_if_missing(DEEPGOZERO_GIT, DEEPGOZERO_DIR)\n# clone_if_missing(EL_EMBED_GIT, EL_EMBED_DIR)\n\n# # 2) try to download deepgo dataset (optional but helpful)\n# download_deepgo_data()\n\n# # 3) choose GO input file: prefer deepgo data's go.owl/go.obo if present, else fall back to CAFA input path\n# candidate_go_obo = None\n# # check deepgo data extracted location\n# deepgo_obo1 = DATA_DIR / \"go.obo\"\n# deepgo_obo2 = DATA_DIR / \"go.owl\"\n# if deepgo_obo1.exists():\n# candidate_go_obo = str(deepgo_obo1)\n# elif deepgo_obo2.exists():\n# candidate_go_obo = str(deepgo_obo2)\n# else:\n# # fallback to the CAFA dataset GO (if available in Kaggle input path)\n# cafa_go = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n# if os.path.exists(cafa_go):\n# candidate_go_obo = cafa_go\n\n# if not candidate_go_obo:\n# log.error(\"No GO OBO/OWL file found in data directory or CAFA path; cannot proceed with EL normalization.\")\n# # but we can still try node2vec fallback if CAFA go exists in input path - try that\n# cafa_go = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n# if os.path.exists(cafa_go):\n# candidate_go_obo = cafa_go\n# else:\n# log.error(\"No GO OBO available; aborting.\")\n# return\n\n# # 4) attempt EL pipeline\n# try:\n# el_ok = try_el_pipeline(go_input_owl=candidate_go_obo)\n# except Exception as e:\n# log.exception(\"Exception while running EL pipeline: %s\", e)\n# el_ok = False\n\n# if el_ok:\n# log.info(\"EL-embeddings generated successfully.\")\n# return\n\n# # 5) fallback: node2vec on GO OBO\n# try:\n# node2vec_out = ROOT / \"go_embeddings_node2vec.pkl\"\n# node2vec_ok = node2vec_fallback(candidate_go_obo, node2vec_out)\n# if node2vec_ok:\n# log.info(\"Node2vec fallback succeeded. Embeddings available at %s\", node2vec_out)\n# else:\n# log.error(\"Node2vec fallback failed.\")\n# except Exception as e:\n# log.exception(\"Node2vec fallback raised exception: %s\", e)\n\n# if __name__ == \"__main__\":\n# main() ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T12:37:23.927991Z","iopub.execute_input":"2026-03-02T12:37:23.928305Z","iopub.status.idle":"2026-03-02T13:11:39.322793Z","shell.execute_reply.started":"2026-03-02T12:37:23.928282Z","shell.execute_reply":"2026-03-02T13:11:39.322128Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # cell: download_and_map_deepgozero_precomputed.py\n# # Run in one Kaggle notebook cell.\n# import os\n# import sys\n# import logging\n# import tarfile\n# import shutil\n# import subprocess\n# from pathlib import Path\n# import urllib.request\n# import pandas as pd\n# from Bio import SeqIO\n\n# # ---------- CONFIG ----------\n# ROOT = Path(\"/kaggle/working/deepgozero_el_embeddings\")\n# ROOT.mkdir(parents=True, exist_ok=True)\n\n# LOGFILE = ROOT / \"pretrained_download.log\"\n\n# # Authors' precomputed file (published on DeepGO hosts)\n# ZERO_PRED_URL = \"https://deepgo.cbrc.kaust.edu.sa/data/deepgozero/zero_predictions.tar.gz\"\n\n# # Optional: if you have a direct URL to a pretrained deepgozero model checkpoint, set it here.\n# # Example: PRETRAINED_MODEL_URL = \"https://example.com/path/to/deepgozero.th\"\n# PRETRAINED_MODEL_URL = None # <-- set if you have a model file URL\n\n# # Your dataset locations (adjust if needed)\n# TEST_FASTA = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n# OUT_DIR = ROOT / \"zero_predictions\"\n# OUT_DIR.mkdir(parents=True, exist_ok=True)\n\n# MAPPED_OUT = ROOT / \"mapped_predictions.csv\"\n\n# # ---------- Logging ----------\n# logging.basicConfig(\n# level=logging.INFO,\n# format=\"%(asctime)s [%(levelname)s] %(message)s\",\n# handlers=[logging.StreamHandler(sys.stdout),\n# logging.FileHandler(LOGFILE)]\n# )\n# log = logging.getLogger(\"deepgozero-pretrained\")\n\n# # ---------- Helper functions ----------\n# def download_url(url, dest_path, chunk=1024*64):\n# dest_path = Path(dest_path)\n# if dest_path.exists():\n# log.info(\"File already exists: %s\", dest_path)\n# return True\n# log.info(\"Downloading %s -> %s\", url, dest_path)\n# try:\n# with urllib.request.urlopen(url) as u, open(dest_path, 'wb') as f:\n# meta = u.info()\n# total = int(meta.get(\"Content-Length\", -1))\n# downloaded = 0\n# while True:\n# chunk_data = u.read(chunk)\n# if not chunk_data:\n# break\n# f.write(chunk_data)\n# downloaded += len(chunk_data)\n# if total > 0:\n# pct = downloaded / total * 100\n# log.info(\"Downloaded %.1f%%\", pct)\n# log.info(\"Download finished: %s\", dest_path)\n# return True\n# except Exception as e:\n# log.exception(\"Failed to download %s : %s\", url, e)\n# return False\n\n# def safe_extract_tar(tar_path, extract_to):\n# tar_path = Path(tar_path)\n# extract_to = Path(extract_to)\n# if not tar_path.exists():\n# log.error(\"Tar file does not exist: %s\", tar_path)\n# return False\n# log.info(\"Extracting %s -> %s\", tar_path, extract_to)\n# try:\n# with tarfile.open(tar_path, \"r:gz\") as tar:\n# tar.extractall(path=str(extract_to))\n# log.info(\"Extraction complete.\")\n# return True\n# except Exception as e:\n# log.exception(\"Extraction failed: %s\", e)\n# return False\n\n# def list_files_recursive(path):\n# files = []\n# for root, _, filenames in os.walk(path):\n# for fn in filenames:\n# files.append(os.path.join(root, fn))\n# return files\n\n# # ---------- Step 1: Download zero_predictions.tar.gz ----------\n# tar_local = OUT_DIR / \"zero_predictions.tar.gz\"\n# ok = download_url(ZERO_PRED_URL, tar_local)\n# if not ok:\n# log.warning(\"Could not download zero_predictions from %s. If you have a local/precached copy, place it at %s and re-run.\", ZERO_PRED_URL, tar_local)\n\n# # ---------- Step 2: Extract if download succeeded ----------\n# if tar_local.exists():\n# extracted_ok = safe_extract_tar(tar_local, OUT_DIR)\n# else:\n# extracted_ok = False\n\n# # ---------- Step 3: Inspect extracted files ----------\n# if extracted_ok:\n# files = list_files_recursive(OUT_DIR)\n# log.info(\"Extracted %d files (showing up to 50):\", len(files))\n# for f in files[:50]:\n# log.info(\" - %s\", f)\n# else:\n# files = []\n# log.info(\"No extracted files to inspect.\")\n\n# # ---------- Step 4: Try to load known file formats from extracted content ----------\n# # We attempt to find a reasonable predictions file: .tsv, .csv, .json, .pkl\n# predictions_df = None\n# candidates = [f for f in files if f.lower().endswith((\".tsv\", \".csv\", \".pkl\", \".json\", \".gz\"))]\n# log.info(\"Prediction file candidates: %s\", candidates[:20])\n\n# def try_load_pred_file(path):\n# path = Path(path)\n# try:\n# if path.suffix == \".pkl\":\n# df = pd.read_pickle(path)\n# return df\n# elif path.suffix in [\".tsv\", \".csv\"]:\n# sep = \"\\t\" if path.suffix == \".tsv\" else \",\"\n# df = pd.read_csv(path, sep=sep)\n# return df\n# elif path.suffix == \".json\":\n# return pd.read_json(path, lines=False)\n# elif path.suffix == \".gz\":\n# # try common gz types: maybe .tsv.gz or .csv.gz\n# name = path.name.lower()\n# if name.endswith(\".tsv.gz\"):\n# return pd.read_csv(path, sep=\"\\t\", compression=\"gzip\")\n# if name.endswith(\".csv.gz\"):\n# return pd.read_csv(path, sep=\",\", compression=\"gzip\")\n# # fallback: try read_pickle\n# try:\n# import gzip, pickle\n# with gzip.open(path, \"rb\") as f:\n# return pickle.load(f)\n# except Exception:\n# log.warning(\"Unknown gz file format: %s\", path)\n# return None\n# else:\n# log.warning(\"Unsupported candidate format: %s\", path)\n# return None\n# except Exception as e:\n# log.exception(\"Failed to load prediction file %s : %s\", path, e)\n# return None\n\n# for c in candidates:\n# df_try = try_load_pred_file(c)\n# if df_try is not None:\n# log.info(\"Loaded predictions candidate: %s (shape=%s)\", c, getattr(df_try, \"shape\", None))\n# predictions_df = df_try\n# pred_file_used = c\n# break\n\n# if predictions_df is None:\n# log.warning(\"No predictions file could be auto-loaded from the extracted contents. You can inspect %s manually.\", OUT_DIR)\n# else:\n# log.info(\"Using predictions file: %s\", pred_file_used)\n# # Show top rows\n# log.info(\"Predictions preview:\\n%s\", predictions_df.head(10).to_string())\n\n# # ---------- Step 5: Load test FASTA (accessions + sequences + parsed metadata) ----------\n# # We'll parse FASTA headers robustly and build a DataFrame with accession as key.\n# def parse_uniprot_header_universal(header):\n# # fallback simple parser that extracts accession-like token and retains remainder\n# # Many CAFA headers have the accession as first token or form sp|ACC|NAME or just ACC\n# # We'll try sp|...|... first then fallback to first token.\n# header = header.strip()\n# if header.startswith(\"sp|\") or header.startswith(\"tr|\"):\n# parts = header.split(\" \", 1)\n# id_part = parts[0]\n# rest = parts[1] if len(parts) > 1 else \"\"\n# toks = id_part.split(\"|\")\n# if len(toks) >= 3:\n# prefix, accession, entry = toks[0], toks[1], toks[2]\n# return accession, entry, rest.strip()\n# else:\n# return id_part, \"\", header\n# # fallback: first token is accession, sometimes header like \"A0A0C5B5G6 9606\"\n# first = header.split()[0]\n# # if first looks like accession (letters+digits) return it\n# return first, \"\", header\n\n# test_records = []\n# if os.path.exists(TEST_FASTA):\n# log.info(\"Parsing test FASTA: %s\", TEST_FASTA)\n# for rec in SeqIO.parse(TEST_FASTA, \"fasta\"):\n# acc, entry, desc = parse_uniprot_header_universal(rec.description)\n# test_records.append({\n# \"accession\": acc,\n# \"entry_name\": entry,\n# \"description_raw\": desc,\n# \"sequence\": str(rec.seq)\n# })\n# df_test = pd.DataFrame(test_records)\n# log.info(\"Loaded test FASTA: %d sequences\", len(df_test))\n# log.info(\"Sample rows:\\n%s\", df_test.head(5).to_string())\n# else:\n# log.error(\"Test FASTA not found at: %s\", TEST_FASTA)\n# df_test = pd.DataFrame([])\n\n# # ---------- Step 6: Merge predictions (if we loaded them) ----------\n# if predictions_df is not None and not df_test.empty:\n# # Attempt 2 common schemas:\n# # - Many prediction dumps have columns like ['protein_id','go_id','score'] or ['Entry','GO','score']\n# # Try to detect columns\n# cols = [c.lower() for c in predictions_df.columns]\n# log.info(\"Detected prediction columns: %s\", cols)\n\n# # Normalize column names (common possibilities)\n# mapping = {}\n# if any(c in cols for c in (\"protein_id\",\"entry\",\"accession\",\"proteins\")):\n# for c in (\"protein_id\",\"Entry\",\"entry\",\"accession\",\"proteins\",\"protein\"):\n# if c in predictions_df.columns:\n# mapping[c] = \"accession\"\n# break\n# if any(c in cols for c in (\"go\",\"go_id\",\"go_term\",\"term\")):\n# for c in (\"go\",\"go_id\",\"go_term\",\"term\",\"GO\"):\n# if c in predictions_df.columns:\n# mapping[c] = \"go_id\"\n# break\n# if any(c in cols for c in (\"score\",\"prob\",\"probability\",\"value\")):\n# for c in (\"score\",\"prob\",\"probability\",\"value\"):\n# if c in predictions_df.columns:\n# mapping[c] = \"score\"\n# break\n\n# if mapping:\n# log.info(\"Renaming prediction columns according to detected mapping: %s\", mapping)\n# predictions_df = predictions_df.rename(columns=mapping)\n# else:\n# log.warning(\"Could not automatically map prediction columns. Available columns: %s\", list(predictions_df.columns))\n# # show and stop merging\n# # If file is other shape (e.g., JSON where keys are proteins mapping to dicts), try other loaders\n# if \"accession\" in predictions_df.columns and \"go_id\" in predictions_df.columns:\n# # some rows can show repeated accessions -> pivot to list of (go,score)\n# preds_grouped = predictions_df.groupby(\"accession\").apply(\n# lambda g: list(zip(g[\"go_id\"].astype(str), g[\"score\"].astype(float) if \"score\" in g.columns else [1.0]*len(g)))\n# ).reset_index().rename(columns={0:\"predicted_terms\"})\n# # merge to df_test on accession\n# df_merged = df_test.merge(preds_grouped, left_on=\"accession\", right_on=\"accession\", how=\"left\")\n# df_merged[\"predicted_terms\"] = df_merged[\"predicted_terms\"].apply(lambda x: x if isinstance(x, list) else [])\n# log.info(\"Merged predictions to test proteins. Result shape: %s\", df_merged.shape)\n# log.info(\"Sample merged rows:\\n%s\", df_merged.head(5).to_string())\n# # Save\n# df_merged.to_csv(MAPPED_OUT, index=False)\n# log.info(\"Saved mapped predictions to: %s\", MAPPED_OUT)\n# else:\n# log.warning(\"Prediction file did not contain expected columns after rename. Please inspect the extracted files at %s\", OUT_DIR)\n# else:\n# log.warning(\"Skipping merge: either no predictions loaded or no test FASTA found.\")\n\n# # ---------- Step 7: (Optional) If user supplied PRETRAINED_MODEL_URL, attempt to download & run prediction ----------\n# if PRETRAINED_MODEL_URL:\n# log.info(\"PRETRAINED_MODEL_URL provided; attempting to download model and run deepgozero_predict.py.\")\n# model_path = ROOT / \"pretrained_deepgozero_model.th\"\n# ok = download_url(PRETRAINED_MODEL_URL, model_path)\n# if ok and model_path.exists():\n# # Clone repo if missing\n# repo_dir = ROOT / \"deepgozero_repo\"\n# if not repo_dir.exists():\n# log.info(\"Cloning DeepGOZero repo for using deepgozero_predict.py\")\n# res = subprocess.run([\"git\", \"clone\", \"--depth\", \"1\", \"https://github.com/bio-ontology-research-group/deepgozero.git\", str(repo_dir)], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n# log.info(res.stdout[:2000])\n# predict_script = repo_dir / \"deepgozero_predict.py\"\n# if predict_script.exists():\n# # Example invocation (you may need to adapt --data-root and data files)\n# cmd = [sys.executable, str(predict_script),\n# \"--data-root\", str(repo_dir / \"data\"),\n# \"--ont\", \"mf\",\n# \"--data-file\", \"swissprot.pkl\",\n# \"--device\", \"cuda:0\"]\n# log.info(\"Would run: %s\", \" \".join(cmd))\n# log.info(\"NOTE: Running deepgozero_predict.py requires the DeepGO data formatting under data/ (interpros, terms, etc.) which is not set up in this notebook. See README in repo for data preparation steps.\")\n# else:\n# log.error(\"deepgozero_predict.py not found in cloned repo at %s\", predict_script)\n# else:\n# log.error(\"Failed to download pretrained model from PRETRAINED_MODEL_URL: %s\", PRETRAINED_MODEL_URL)\n\n# log.info(\"Script finished. Inspect files under %s and the notebook log for details.\", ROOT)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import shutil\n\nshutil.rmtree(\"/kaggle/working/deepgozero_el_embeddings\", ignore_errors=True)\nprint(\"Deleted folder deepgozero_el_embeddings\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T13:59:35.913711Z","iopub.execute_input":"2026-03-02T13:59:35.914397Z","iopub.status.idle":"2026-03-02T13:59:36.980629Z","shell.execute_reply.started":"2026-03-02T13:59:35.914367Z","shell.execute_reply":"2026-03-02T13:59:36.979659Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport zipfile\n\nembed_dir = \"/kaggle/working/protbert_embeddings\"\nzip_path = \"/kaggle/working/train_protbert_embeddings.zip\"\n\nwith zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:\n for fname in os.listdir(embed_dir):\n if fname.startswith(\"train_protbert_embeddings_\") and fname.endswith(\".pt\"):\n zipf.write(os.path.join(embed_dir, fname), fname)\n\nprint(\"Zipped file saved at:\", zip_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T14:00:30.124847Z","iopub.execute_input":"2026-03-02T14:00:30.125382Z","iopub.status.idle":"2026-03-02T14:03:09.644132Z","shell.execute_reply.started":"2026-03-02T14:00:30.125358Z","shell.execute_reply":"2026-03-02T14:03:09.643369Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install node2vec==0.5.0 ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T15:15:26.886635Z","iopub.execute_input":"2026-03-02T15:15:26.887349Z","iopub.status.idle":"2026-03-02T15:15:28.822998Z","shell.execute_reply.started":"2026-03-02T15:15:26.887316Z","shell.execute_reply":"2026-03-02T15:15:28.822034Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install smart_open==1.9.0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T15:18:54.156906Z","iopub.execute_input":"2026-03-02T15:18:54.157477Z","iopub.status.idle":"2026-03-02T15:19:00.968735Z","shell.execute_reply.started":"2026-03-02T15:18:54.157451Z","shell.execute_reply":"2026-03-02T15:19:00.967964Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pickle\nimport numpy as np\nimport pandas as pd\nimport networkx as nx\nfrom tqdm import tqdm\nfrom node2vec import Node2Vec\nimport obonet\n\n# ---------- Inputs ----------\nGO_OBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nTRAIN_TERMS_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n\n# Where to save embeddings\nOUT_EMB_DIR = \"/kaggle/working/goa_node2vec_embeddings\"\nos.makedirs(OUT_EMB_DIR, exist_ok=True)\n\n# Node2vec params (tune these for memory/time)\nDIMENSIONS = 200\nWALK_LENGTH = 30\nNUM_WALKS = 20\nWINDOW_SIZE = 10\nWORKERS = 4\nEPOCHS = 5\n\nprint(\"Loading GO ontology...\")\n\n# ---------- Build GO and GOA graph ----------\nprint(\"Parsing GO OBO...\")\ngo_graph = obonet.read_obo(GO_OBO_PATH)\n\nprint(\"Converting GO edges to undirected NetworkX graph...\")\nG = nx.Graph()\n\n# Step 1: add GO term nodes + term-term edges (ontology structure)\nfor u, v, _ in go_graph.edges(data=True):\n G.add_node(u) # GO term\n G.add_node(v)\n G.add_edge(u, v)\n\nprint(f\"Added GO term graph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges\")\n\nprint(\"Loading protein-term annotations...\")\ndf_terms = pd.read_csv(TRAIN_TERMS_PATH, sep=\"\\t\", names=[\"entry\", \"go_term\", \"aspect\"])\n\nprint(\"Adding protein nodes and annotation edges...\")\nfor idx, row in tqdm(df_terms.iterrows(), total=len(df_terms)):\n prot = str(row[\"entry\"])\n term = str(row[\"go_term\"])\n # add protein and annotation edge\n G.add_node(prot)\n G.add_edge(prot, term)\n\nprint(f\"Combined GOA graph: nodes={G.number_of_nodes()} edges={G.number_of_edges()}\")\n\n# ---------- Run node2vec ----------\nprint(\"Initializing Node2Vec random walks...\")\nnode2vec = Node2Vec(G, dimensions=DIMENSIONS, walk_length=WALK_LENGTH,\n num_walks=NUM_WALKS, workers=WORKERS)\n\nprint(\"Fitting node2vec (Word2Vec)...\")\nmodel = node2vec.fit(window=WINDOW_SIZE, min_count=1, batch_words=4, epochs=EPOCHS)\n\n# ---------- Extract embeddings ----------\nprint(\"Extracting embeddings for GO terms and proteins...\")\n\nembeddings = {}\ngo_term_embs = {}\nprotein_embs = {}\n\nfor node in model.wv.index_to_key:\n vec = model.wv[node]\n embeddings[node] = vec\n # GO terms and proteins can be separated by ID format\n if node.startswith(\"GO:\"):\n go_term_embs[node] = vec\n else:\n protein_embs[node] = vec\n\nprint(f\"Total embedded nodes: {len(embeddings)}\")\nprint(f\"GO terms: {len(go_term_embs)}; proteins: {len(protein_embs)}\")\n\n# ---------- Save embeddings ----------\ngo_out = os.path.join(OUT_EMB_DIR, \"go_term_node2vec_emb.pkl\")\nprot_out = os.path.join(OUT_EMB_DIR, \"protein_node2vec_emb.pkl\")\nall_out = os.path.join(OUT_EMB_DIR, \"all_node2vec_emb.pkl\")\n\nwith open(go_out, \"wb\") as f:\n pickle.dump(go_term_embs, f)\nprint(f\"GO term embeddings saved to: {go_out}\")\n\nwith open(prot_out, \"wb\") as f:\n pickle.dump(protein_embs, f)\nprint(f\"Protein embeddings saved to: {prot_out}\")\n\nwith open(all_out, \"wb\") as f:\n pickle.dump(embeddings, f)\nprint(f\"All node2vec embeddings saved to: {all_out}\")\n\nprint(\"Node2Vec GOA embedding generation complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T15:19:17.021099Z","iopub.execute_input":"2026-03-02T15:19:17.021382Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install torch_geometric torch-cluster","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T17:02:23.072321Z","iopub.execute_input":"2026-03-02T17:02:23.073055Z","iopub.status.idle":"2026-03-02T17:02:26.712132Z","shell.execute_reply.started":"2026-03-02T17:02:23.073016Z","shell.execute_reply":"2026-03-02T17:02:26.711323Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\n\ntorch_version = torch.__version__.split('+')[0]\ncuda_version = torch.version.cuda\n\nif cuda_version:\n # Formats '12.1' to 'cu121'\n cuda_formatted = f\"cu{cuda_version.replace('.', '')}\" \nelse:\n cuda_formatted = \"cpu\"\n\nprint(f\"Your PyTorch version: {torch_version}\")\nprint(f\"Your CUDA version: {cuda_formatted}\")\nprint(f\"Wheel URL suffix: torch-{torch_version}+{cuda_formatted}.html\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T17:04:44.084043Z","iopub.execute_input":"2026-03-02T17:04:44.084986Z","iopub.status.idle":"2026-03-02T17:04:44.091207Z","shell.execute_reply.started":"2026-03-02T17:04:44.08495Z","shell.execute_reply":"2026-03-02T17:04:44.090342Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install torch-cluster -f https://data.pyg.org/whl/torch-2.6.0+cu124.html","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T17:05:22.892248Z","iopub.execute_input":"2026-03-02T17:05:22.892602Z","iopub.status.idle":"2026-03-02T17:05:25.433295Z","shell.execute_reply.started":"2026-03-02T17:05:22.892574Z","shell.execute_reply":"2026-03-02T17:05:25.432483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\n\n# Get exact versions\ntorch_version = torch.__version__.split('+')[0]\ncuda_version = torch.version.cuda\ncuda_formatted = f\"cu{cuda_version.replace('.', '')}\" if cuda_version else \"cpu\"\n\nwhl_url = f\"https://data.pyg.org/whl/torch-{torch_version}+{cuda_formatted}.html\"\nprint(f\"Fetching wheels from: {whl_url}\")\n\n# Force install both pyg-lib and torch-cluster\n!pip install pyg-lib torch-cluster torch-scatter torch-sparse -f {whl_url}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T17:07:35.691734Z","iopub.execute_input":"2026-03-02T17:07:35.69244Z","iopub.status.idle":"2026-03-02T17:07:39.140718Z","shell.execute_reply.started":"2026-03-02T17:07:35.692413Z","shell.execute_reply":"2026-03-02T17:07:39.139802Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pickle\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport obonet\n\nimport torch\nimport torch.nn as nn\nfrom torch_geometric.nn import Node2Vec\n\n# ---------- Inputs ----------\nGO_OBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nTRAIN_TERMS_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\nOUT_EMB_DIR = \"/kaggle/working/goa_node2vec_embeddings\"\nos.makedirs(OUT_EMB_DIR, exist_ok=True)\n\n# Node2vec params\nDIMENSIONS = 200\nWALK_LENGTH = 30\nNUM_WALKS = 20\nWINDOW_SIZE = 10\nWORKERS = 4\nEPOCHS = 5\nBATCH_SIZE = 256 # Added for PyTorch DataLoader\n\nprint(\"Loading GO ontology...\")\n\n# ---------- Build PyTorch Geometric Graph ----------\n# We use integer IDs for PyTorch, so we need mapping dictionaries\nnode2idx = {}\nidx2node = []\n\ndef get_node_idx(node_name):\n if node_name not in node2idx:\n node2idx[node_name] = len(idx2node)\n idx2node.append(node_name)\n return node2idx[node_name]\n\nprint(\"Parsing GO OBO...\")\ngo_graph = obonet.read_obo(GO_OBO_PATH)\n\nsrc, dst = [], []\n\nprint(\"Extracting GO edges...\")\nfor u, v, _ in go_graph.edges(data=True):\n uid = get_node_idx(u)\n vid = get_node_idx(v)\n # Add undirected edges\n src.extend([uid, vid])\n dst.extend([vid, uid])\n\nprint(\"Loading protein-term annotations...\")\n# Load faster by dropping missing values if any exist\ndf_terms = pd.read_csv(TRAIN_TERMS_PATH, sep=\"\\t\", names=[\"entry\", \"go_term\", \"aspect\"]).dropna()\n\nprint(\"Adding protein nodes and annotation edges...\")\nfor row in tqdm(df_terms.itertuples(index=False), total=len(df_terms)):\n prot = str(row.entry)\n term = str(row.go_term)\n \n uid = get_node_idx(prot)\n vid = get_node_idx(term)\n \n # Add undirected edges\n src.extend([uid, vid])\n dst.extend([vid, uid])\n\nnum_nodes = len(idx2node)\nprint(f\"Total Unique Nodes: {num_nodes}\")\nprint(f\"Total Edges (Undirected): {len(src)}\")\n\n# Create PyTorch Edge Index Tensor\nedge_index = torch.tensor([src, dst], dtype=torch.long)\n\n# ---------- Setup PyTorch Node2Vec Model ----------\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f\"Initializing Node2Vec on {device}...\")\nmodel = Node2Vec(\n edge_index, \n embedding_dim=DIMENSIONS, \n walk_length=WALK_LENGTH,\n context_size=WINDOW_SIZE, \n walks_per_node=NUM_WALKS,\n num_negative_samples=1, \n p=1, q=1, sparse=True\n)\n\n# Wrapper to make Node2Vec compatible with nn.DataParallel's forward pass\nclass ParallelNode2Vec(nn.Module):\n def __init__(self, base_model):\n super().__init__()\n self.model = base_model\n\n def forward(self, pos_rw, neg_rw):\n return self.model.loss(pos_rw, neg_rw)\n\n# Wrap in DataParallel if multiple GPUs are available\nif torch.cuda.device_count() > 1:\n print(f\"Using {torch.cuda.device_count()} GPUs with DataParallel!\")\n dp_model = nn.DataParallel(ParallelNode2Vec(model))\nelse:\n dp_model = ParallelNode2Vec(model)\n\ndp_model = dp_model.to(device)\n\n# PyG Node2Vec Loader uses multiprocessing to generate walks fast on CPU\nloader = model.loader(batch_size=BATCH_SIZE, shuffle=True, num_workers=WORKERS)\n\n# Using SparseAdam because Node2Vec uses Sparse Tensors for embeddings\noptimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01)\n\n# ---------- Train Node2Vec ----------\nprint(\"Starting Training...\")\nfor epoch in range(1, EPOCHS + 1):\n dp_model.train()\n total_loss = 0\n \n for pos_rw, neg_rw in tqdm(loader, desc=f\"Epoch {epoch}/{EPOCHS}\"):\n optimizer.zero_grad()\n \n # Calculate loss (Parallelized across GPUs if DataParallel is active)\n loss = dp_model(pos_rw.to(device), neg_rw.to(device)).mean()\n \n loss.backward()\n optimizer.step()\n total_loss += loss.item()\n \n print(f\"Epoch: {epoch:02d}, Loss: {total_loss / len(loader):.4f}\")\n\n# ---------- Extract Embeddings ----------\nprint(\"Extracting embeddings for GO terms and proteins...\")\n\nmodel.eval()\nwith torch.no_grad():\n # Move embeddings to CPU and convert to numpy\n all_embeddings = model.embedding.weight.cpu().numpy()\n\nembeddings = {}\ngo_term_embs = {}\nprotein_embs = {}\n\nfor idx, node_name in enumerate(idx2node):\n vec = all_embeddings[idx]\n embeddings[node_name] = vec\n \n if node_name.startswith(\"GO:\"):\n go_term_embs[node_name] = vec\n else:\n protein_embs[node_name] = vec\n\nprint(f\"Total embedded nodes: {len(embeddings)}\")\nprint(f\"GO terms: {len(go_term_embs)}; proteins: {len(protein_embs)}\")\n\n# ---------- Save Embeddings ----------\ngo_out = os.path.join(OUT_EMB_DIR, \"go_term_node2vec_emb.pkl\")\nprot_out = os.path.join(OUT_EMB_DIR, \"protein_node2vec_emb.pkl\")\nall_out = os.path.join(OUT_EMB_DIR, \"all_node2vec_emb.pkl\")\n\nwith open(go_out, \"wb\") as f:\n pickle.dump(go_term_embs, f)\nprint(f\"GO term embeddings saved to: {go_out}\")\n\nwith open(prot_out, \"wb\") as f:\n pickle.dump(protein_embs, f)\nprint(f\"Protein embeddings saved to: {prot_out}\")\n\nwith open(all_out, \"wb\") as f:\n pickle.dump(embeddings, f)\nprint(f\"All node2vec embeddings saved to: {all_out}\")\n\nprint(\"Node2Vec GOA embedding generation complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T17:07:49.382948Z","iopub.execute_input":"2026-03-02T17:07:49.383291Z","iopub.status.idle":"2026-03-02T17:07:57.012163Z","shell.execute_reply.started":"2026-03-02T17:07:49.383261Z","shell.execute_reply":"2026-03-02T17:07:57.011133Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from IPython.display import FileLink\n\n# Replace with the actual name of your zip file\nFileLink('/kaggle/working/train_protbert_embeddings.zip')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-02T17:10:32.676729Z","iopub.execute_input":"2026-03-02T17:10:32.67758Z","iopub.status.idle":"2026-03-02T17:10:32.682626Z","shell.execute_reply.started":"2026-03-02T17:10:32.677545Z","shell.execute_reply":"2026-03-02T17:10:32.681962Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---","metadata":{}},{"cell_type":"code","source":"TRAIN_DIR = '/kaggle/input/cafa-6-protein-function-prediction/Train'\nTEST_DIR = '/kaggle/input/cafa-6-protein-function-prediction/Test'\nDATA_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\nEMBEDDINGS_DIR = '/kaggle/working/embeddings'\nMODEL_DIR = '/kaggle/working/models'\n\n# Create directories\nos.makedirs(EMBEDDINGS_DIR, exist_ok=True)\nos.makedirs(MODEL_DIR, exist_ok=True)\n\n# File paths\nTRAIN_FASTA = f'{TRAIN_DIR}/train_sequences.fasta'\nTRAIN_TERMS = f'{TRAIN_DIR}/train_terms.tsv'\nTEST_FASTA = f'{TEST_DIR}/testsuperset.fasta'\nIA_FILE = f'{DATA_DIR}/IA.tsv'\n\nprint(\"āœ“ Paths configured\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_fasta(fasta_file):\n \"\"\"Parse FASTA file.\"\"\"\n sequences = {}\n current_id = None\n current_seq = []\n \n with open(fasta_file, 'r') as f:\n for line in f:\n line = line.strip()\n if line.startswith('>'):\n if current_id:\n sequences[current_id] = ''.join(current_seq)\n header = line[1:]\n current_id = header.split('|')[1] if '|' in header else header.split()[0]\n current_seq = []\n else:\n current_seq.append(line)\n if current_id:\n sequences[current_id] = ''.join(current_seq)\n return sequences\n\n# Load sequences\nprint(\"Loading training sequences...\")\ntrain_sequences = parse_fasta(TRAIN_FASTA)\nprint(f\" Loaded {len(train_sequences)} training sequences\")\n\nprint(\"Loading test sequences...\")\ntest_sequences = parse_fasta(TEST_FASTA)\nprint(f\" Loaded {len(test_sequences)} test sequences\")\n\n# Load GO terms\nprint(\"Loading GO terms...\")\ndf_terms = pd.read_csv(TRAIN_TERMS, sep='\\t', header=None, \n names=['protein_id', 'go_term', 'ontology'])\nprotein_terms = df_terms.groupby('protein_id')['go_term'].apply(list).to_dict()\nall_go_terms = sorted(df_terms['go_term'].unique())\nprint(f\" Loaded {len(protein_terms)} annotated proteins\")\nprint(f\" Total unique GO terms: {len(all_go_terms)}\")\n\nprint(\"\\nāœ“ Data loaded successfully\")\n\n# Quick data exploration\nprint(\"\\nšŸ“Š Data Summary:\")\nprint(f\" Average sequence length: {np.mean([len(s) for s in list(train_sequences.values())[:1000]]):.0f}\")\nprint(f\" Average GO terms per protein: {df_terms.groupby('protein_id').size().mean():.1f}\")\nprint(f\" Example protein: {list(train_sequences.keys())[0]}\")\nprint(f\" Example sequence (first 50 chars): {list(train_sequences.values())[0][:50]}...\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_EMB_FILE = f'{EMBEDDINGS_DIR}/train_embeddings.pkl'\n\nif os.path.exists(TRAIN_EMB_FILE):\n print(\"Loading cached training embeddings...\")\n with open(TRAIN_EMB_FILE, 'rb') as f:\n train_embeddings = pickle.load(f)\n print(f\"āœ“ Loaded {len(train_embeddings)} training embeddings from cache\")\nelse:\n print(\"Extracting ESM-2 embeddings for training set...\")\n print(\"āš ļø This will take 2-6 hours. Get some coffee! ā˜•\")\n \n # Load ESM-2 model\n MODEL_NAME = \"facebook/esm2_t33_650M_UR50D\"\n # Alternative: \"facebook/esm2_t30_150M_UR50D\" (smaller, faster)\n \n print(f\"Loading model: {MODEL_NAME}\")\n tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n model = EsmModel.from_pretrained(MODEL_NAME)\n \n device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n model = model.to(device)\n model.eval()\n print(f\"Model loaded on: {device}\")\n \n # Extract embeddings\n train_embeddings = {}\n protein_ids = list(train_sequences.keys())\n batch_size = 4 # Reduce if you get memory errors\n \n with torch.no_grad():\n for i in tqdm(range(0, len(protein_ids), batch_size)):\n batch_ids = protein_ids[i:i+batch_size]\n batch_seqs = [train_sequences[pid] for pid in batch_ids]\n \n # Tokenize\n inputs = tokenizer(batch_seqs, return_tensors='pt', \n padding=True, truncation=True, \n max_length=1024)\n inputs = {k: v.to(device) for k, v in inputs.items()}\n \n # Extract embeddings\n outputs = model(**inputs)\n \n # Mean pooling\n embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()\n \n for j, pid in enumerate(batch_ids):\n train_embeddings[pid] = embeddings[j]\n \n # Save embeddings\n print(f\"\\nSaving embeddings to {TRAIN_EMB_FILE}\")\n with open(TRAIN_EMB_FILE, 'wb') as f:\n pickle.dump(train_embeddings, f)\n \n print(f\"āœ“ Extracted and saved {len(train_embeddings)} embeddings\")\n\nprint(f\"\\nāœ“ Training embeddings ready\")\nprint(f\" Embedding dimension: {list(train_embeddings.values())[0].shape[0]}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pickle\n\nEMB_PATH = \"/kaggle/input/cafa6-files/train_embeddings.pkl\"\n\nwith open(EMB_PATH, \"rb\") as f:\n train_embeddings = pickle.load(f)\n\nprint(f\"Loaded {len(train_embeddings)} embeddings successfully āœ…\")\nprint(f\"Embedding dimension: {list(train_embeddings.values())[0].shape}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install fair-esm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Robust ESM embedding extraction with checkpointing, resume, and memory-safety\n# Paste this into your Kaggle notebook (requires fair-esm: pip install fair-esm)\n\nimport os, pickle, gc, time\nfrom tqdm.auto import tqdm\nimport torch\nimport esm\n\n# === User variables ===\nEMBEDDINGS_DIR = globals().get(\"EMBEDDINGS_DIR\", \"/kaggle/working/embeddings\")\nos.makedirs(EMBEDDINGS_DIR, exist_ok=True)\nTRAIN_EMB_FILE = f'{EMBEDDINGS_DIR}/train_embeddings.pkl'\nCHECKPOINT_FILE = f'{EMBEDDINGS_DIR}/train_embeddings_checkpoint.pkl'\n\n# Model + performance knobs\nMODEL_NAME = \"esm2_t30_150M_UR50D\"\nUSE_HALF = True\nBATCH_SIZE = 4 # Increased since using multiple GPUs\nCPU_BATCH_FOR_LONG = 1\nLONG_SEQ_THRESHOLD = 1000\nCHECKPOINT_EVERY = 500\nMAX_SEQ_LENGTH = 1022\n# =======================\n\nif 'train_sequences' not in globals():\n raise RuntimeError(\"Please define `train_sequences` (dict mapping protein_id -> sequence) before running this cell.\")\n\n# Multi-GPU setup\nnum_gpus = torch.cuda.device_count()\nif num_gpus > 1:\n print(f\"āœ“ Found {num_gpus} GPUs - using DataParallel\")\n device = torch.device(\"cuda\")\n USE_MULTI_GPU = True\nelif num_gpus == 1:\n print(\"āœ“ Using single GPU\")\n device = torch.device(\"cuda\")\n USE_MULTI_GPU = False\nelse:\n print(\"āš ļø No GPU found - using CPU\")\n device = torch.device(\"cpu\")\n USE_MULTI_GPU = False\n\n# small safety: try reduce CUDA fragmentation\nos.environ.setdefault(\"PYTORCH_CUDA_ALLOC_CONF\", \"expandable_segments:True\")\n\n# load or resume embeddings\nif os.path.exists(TRAIN_EMB_FILE):\n print(\"Loading cached training embeddings...\")\n with open(TRAIN_EMB_FILE, 'rb') as f:\n train_embeddings = pickle.load(f)\n print(f\"āœ“ Loaded {len(train_embeddings)} training embeddings from cache\")\nelse:\n if os.path.exists(CHECKPOINT_FILE):\n print(\"Resuming from checkpoint...\")\n with open(CHECKPOINT_FILE, 'rb') as f:\n train_embeddings = pickle.load(f)\n print(f\"āœ“ Loaded checkpoint with {len(train_embeddings)} embeddings\")\n else:\n train_embeddings = {}\n\n print(\"Extracting ESM-2 embeddings for training set...\")\n print(\"āš ļø This will take 2-6 hours. Get some coffee! ā˜•\")\n print(f\"Loading model: {MODEL_NAME}\")\n\n # Load model\n try:\n model, alphabet = esm.pretrained.load_model_and_alphabet(MODEL_NAME)\n except Exception as e:\n print(f\"Error loading {MODEL_NAME}: {e}\")\n print(\"Trying alternative loading method...\")\n try:\n model, alphabet = getattr(esm.pretrained, MODEL_NAME)()\n except Exception as e2:\n print(f\"Failed to load model: {e2}\")\n raise\n\n batch_converter = alphabet.get_batch_converter()\n\n # Wrap model with DataParallel if multiple GPUs\n if USE_MULTI_GPU:\n model = torch.nn.DataParallel(model)\n print(f\"āœ“ Model wrapped with DataParallel across {num_gpus} GPUs\")\n \n model = model.to(device)\n model.eval()\n \n if device.type == 'cuda' and USE_HALF:\n try:\n model.half()\n print(\"Converted model to float16 to save memory.\")\n except Exception as e:\n print(\"Warning: model.half() failed:\", e)\n\n print(f\"Model loaded on: {device}\")\n\n # Build list of IDs to process\n all_ids = list(train_sequences.keys())\n pending_ids = [pid for pid in all_ids if pid not in train_embeddings]\n total_pending = len(pending_ids)\n print(f\"Total sequences: {len(all_ids)}; Remaining to process: {total_pending}\")\n\n short_ids = [pid for pid in pending_ids if len(train_sequences[pid]) < LONG_SEQ_THRESHOLD]\n long_ids = [pid for pid in pending_ids if len(train_sequences[pid]) >= LONG_SEQ_THRESHOLD]\n print(f\" Short (<{LONG_SEQ_THRESHOLD}): {len(short_ids)} Long (ā‰„{LONG_SEQ_THRESHOLD}): {len(long_ids)}\")\n\n def process_id_list(id_list, device_target, initial_bs):\n idx = 0\n N = len(id_list)\n bs = initial_bs\n oom_count = 0\n\n while idx < N:\n batch_ids = id_list[idx: idx + bs]\n batch_data = [(pid, train_sequences[pid][:MAX_SEQ_LENGTH]) for pid in batch_ids]\n\n try:\n _, _, tokens = batch_converter(batch_data)\n tokens = tokens.to(device_target)\n\n if device_target.type == 'cuda' and USE_HALF:\n context = torch.amp.autocast(device_type='cuda', dtype=torch.float16)\n else:\n class _null_ctx:\n def __enter__(self): return None\n def __exit__(self, exc_type, exc, tb): return False\n context = _null_ctx()\n\n with torch.no_grad():\n with context:\n # Access correct model for DataParallel\n num_layers = model.module.num_layers if USE_MULTI_GPU else model.num_layers\n results = model(tokens, repr_layers=[num_layers], return_contacts=False)\n\n rep_key = list(results[\"representations\"].keys())[-1]\n token_representations = results[\"representations\"][rep_key]\n embeddings = token_representations.mean(dim=1).detach().cpu().float().numpy()\n\n for j, pid in enumerate(batch_ids):\n train_embeddings[pid] = embeddings[j]\n\n idx += bs\n oom_count = 0\n\n if (len(train_embeddings) % CHECKPOINT_EVERY) == 0:\n print(f\"Checkpoint: saved {len(train_embeddings)} embeddings so far (idx {idx}/{N} in this partition).\")\n with open(CHECKPOINT_FILE, 'wb') as f:\n pickle.dump(train_embeddings, f)\n\n del tokens, results, token_representations, embeddings\n if device_target.type == 'cuda':\n torch.cuda.empty_cache()\n gc.collect()\n\n except RuntimeError as e:\n msg = str(e).lower()\n if 'out of memory' in msg or 'cuda out of memory' in msg:\n oom_count += 1\n old_bs = bs\n bs = max(1, bs // 2)\n print(f\"CUDA OOM (batch {old_bs}). Reducing batch size -> {bs} and retrying batch.\")\n \n if 'tokens' in locals():\n del tokens\n if device_target.type == 'cuda':\n torch.cuda.empty_cache()\n gc.collect()\n \n if bs == 1 and device_target.type == 'cuda' and oom_count >= 2:\n print(\"āš ļø Persistent OOM at batch_size=1. Switching to CPU for remaining sequences.\")\n device_target = torch.device('cpu')\n # Unwrap DataParallel before moving to CPU\n if USE_MULTI_GPU:\n model_cpu = model.module.cpu().float()\n else:\n model_cpu = model.cpu().float()\n torch.cuda.empty_cache()\n gc.collect()\n \n continue\n else:\n raise\n\n except KeyboardInterrupt:\n print(\"KeyboardInterrupt received ā€” saving checkpoint before exiting...\")\n with open(CHECKPOINT_FILE, 'wb') as f:\n pickle.dump(train_embeddings, f)\n raise\n\n try:\n if device.type == 'cuda' and len(short_ids) > 0:\n short_ids.sort(key=lambda pid: len(train_sequences[pid]))\n process_id_list(short_ids, device, BATCH_SIZE)\n elif len(short_ids) > 0:\n process_id_list(short_ids, torch.device('cpu'), min(BATCH_SIZE, CPU_BATCH_FOR_LONG))\n\n if len(long_ids) > 0:\n print(\"Processing long sequences on CPU (safer for memory).\")\n if device.type == 'cuda':\n print(\"Moving model to CPU for long sequences...\")\n # Unwrap DataParallel before moving to CPU\n if USE_MULTI_GPU:\n model_cpu = model.module.cpu().float()\n else:\n model_cpu = model.cpu().float()\n torch.cuda.empty_cache()\n gc.collect()\n \n long_ids.sort(key=lambda pid: len(train_sequences[pid]))\n process_id_list(long_ids, torch.device('cpu'), CPU_BATCH_FOR_LONG)\n\n except RuntimeError as e:\n print(f\"RuntimeError during processing: {e}\")\n print(\"Saving checkpoint...\")\n with open(CHECKPOINT_FILE, 'wb') as f:\n pickle.dump(train_embeddings, f)\n print(\"āœ“ Checkpoint saved. You can resume by re-running this cell.\")\n raise\n\n print(f\"\\nSaving embeddings to {TRAIN_EMB_FILE}\")\n with open(TRAIN_EMB_FILE, 'wb') as f:\n pickle.dump(train_embeddings, f)\n \n if os.path.exists(CHECKPOINT_FILE):\n os.remove(CHECKPOINT_FILE)\n\n print(f\"āœ“ Extracted and saved {len(train_embeddings)} embeddings\")\n\nprint(f\"\\nāœ“ Training embeddings ready\")\nprint(f\" Embedding dimension: {list(train_embeddings.values())[0].shape[0]}\")\nprint(f\" Total proteins: {len(train_embeddings)}\")\nprint(\"hello world!!do yo\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Robust ESM embedding extraction with multi-GPU support via data splitting\n# Paste this into your Kaggle notebook (requires fair-esm: pip install fair-esm)\n\nimport os, pickle, gc, time\nfrom tqdm.auto import tqdm\nimport torch\nimport esm\nfrom threading import Thread\nfrom queue import Queue\n\n# === User variables ===\nEMBEDDINGS_DIR = globals().get(\"EMBEDDINGS_DIR\", \"/kaggle/working/embeddings\")\nos.makedirs(EMBEDDINGS_DIR, exist_ok=True)\nTRAIN_EMB_FILE = f'{EMBEDDINGS_DIR}/train_embeddings.pkl'\nCHECKPOINT_FILE = f'{EMBEDDINGS_DIR}/train_embeddings_checkpoint.pkl'\n\n# Model + performance knobs\nMODEL_NAME = \"esm2_t30_150M_UR50D\"\nUSE_HALF = True\nBATCH_SIZE = 4\nCPU_BATCH_FOR_LONG = 1\nLONG_SEQ_THRESHOLD = 1000\nCHECKPOINT_EVERY = 500\nMAX_SEQ_LENGTH = 1022\n# =======================\n\nif 'train_sequences' not in globals():\n raise RuntimeError(\"Please define `train_sequences` (dict mapping protein_id -> sequence) before running this cell.\")\n\n# Multi-GPU setup\nnum_gpus = torch.cuda.device_count()\nprint(f\"āœ“ Found {num_gpus} GPU(s)\")\n\n# small safety: try reduce CUDA fragmentation\nos.environ.setdefault(\"PYTORCH_CUDA_ALLOC_CONF\", \"expandable_segments:True\")\n\n# load or resume embeddings\nif os.path.exists(TRAIN_EMB_FILE):\n print(\"Loading cached training embeddings...\")\n with open(TRAIN_EMB_FILE, 'rb') as f:\n train_embeddings = pickle.load(f)\n print(f\"āœ“ Loaded {len(train_embeddings)} training embeddings from cache\")\nelse:\n if os.path.exists(CHECKPOINT_FILE):\n print(\"Resuming from checkpoint...\")\n with open(CHECKPOINT_FILE, 'rb') as f:\n train_embeddings = pickle.load(f)\n print(f\"āœ“ Loaded checkpoint with {len(train_embeddings)} embeddings\")\n else:\n train_embeddings = {}\n\n print(\"Extracting ESM-2 embeddings for training set...\")\n print(f\"Loading model: {MODEL_NAME}\")\n\n # Build list of IDs to process\n all_ids = list(train_sequences.keys())\n pending_ids = [pid for pid in all_ids if pid not in train_embeddings]\n total_pending = len(pending_ids)\n print(f\"Total sequences: {len(all_ids)}; Remaining to process: {total_pending}\")\n\n short_ids = [pid for pid in pending_ids if len(train_sequences[pid]) < LONG_SEQ_THRESHOLD]\n long_ids = [pid for pid in pending_ids if len(train_sequences[pid]) >= LONG_SEQ_THRESHOLD]\n print(f\" Short (<{LONG_SEQ_THRESHOLD}): {len(short_ids)} Long (ā‰„{LONG_SEQ_THRESHOLD}): {len(long_ids)}\")\n\n def process_on_gpu(gpu_id, id_list, embeddings_dict, lock, checkpoint_file_gpu):\n \"\"\"Process sequences on a specific GPU.\"\"\"\n device = torch.device(f'cuda:{gpu_id}')\n print(f\"[GPU {gpu_id}] Starting with {len(id_list)} sequences\")\n \n # Load model on this GPU\n model, alphabet = esm.pretrained.load_model_and_alphabet(MODEL_NAME)\n batch_converter = alphabet.get_batch_converter()\n \n model = model.to(device)\n model.eval()\n \n if USE_HALF:\n try:\n model.half()\n print(f\"[GPU {gpu_id}] Converted to float16\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Warning: half() failed: {e}\")\n \n idx = 0\n N = len(id_list)\n bs = BATCH_SIZE\n oom_count = 0\n processed_count = 0\n \n # Sort by length\n id_list.sort(key=lambda pid: len(train_sequences[pid]))\n \n while idx < N:\n batch_ids = id_list[idx: idx + bs]\n batch_data = [(pid, train_sequences[pid][:MAX_SEQ_LENGTH]) for pid in batch_ids]\n \n try:\n _, _, tokens = batch_converter(batch_data)\n tokens = tokens.to(device)\n \n if USE_HALF:\n context = torch.amp.autocast(device_type='cuda', dtype=torch.float16)\n else:\n class _null_ctx:\n def __enter__(self): return None\n def __exit__(self, exc_type, exc, tb): return False\n context = _null_ctx()\n \n with torch.no_grad():\n with context:\n results = model(tokens, repr_layers=[model.num_layers], return_contacts=False)\n \n rep_key = list(results[\"representations\"].keys())[-1]\n token_representations = results[\"representations\"][rep_key]\n embeddings = token_representations.mean(dim=1).detach().cpu().float().numpy()\n \n # Thread-safe update\n with lock:\n for j, pid in enumerate(batch_ids):\n embeddings_dict[pid] = embeddings[j]\n processed_count += len(batch_ids)\n \n # Checkpoint periodically\n if processed_count % CHECKPOINT_EVERY == 0:\n print(f\"[GPU {gpu_id}] Processed {processed_count}/{N} sequences\")\n with open(checkpoint_file_gpu, 'wb') as f:\n pickle.dump({k: embeddings_dict[k] for k in id_list if k in embeddings_dict}, f)\n \n idx += bs\n oom_count = 0\n \n # Cleanup\n del tokens, results, token_representations, embeddings\n torch.cuda.empty_cache()\n gc.collect()\n \n except RuntimeError as e:\n msg = str(e).lower()\n if 'out of memory' in msg:\n oom_count += 1\n old_bs = bs\n bs = max(1, bs // 2)\n print(f\"[GPU {gpu_id}] OOM (batch {old_bs}). Reducing to {bs}\")\n \n if 'tokens' in locals():\n del tokens\n torch.cuda.empty_cache()\n gc.collect()\n \n if bs == 1 and oom_count >= 2:\n print(f\"[GPU {gpu_id}] Persistent OOM. Switching to CPU for remaining.\")\n device = torch.device('cpu')\n model = model.cpu().float()\n torch.cuda.empty_cache()\n gc.collect()\n \n continue\n else:\n raise\n \n except KeyboardInterrupt:\n print(f\"[GPU {gpu_id}] Interrupted - saving checkpoint...\")\n with lock:\n with open(checkpoint_file_gpu, 'wb') as f:\n pickle.dump({k: embeddings_dict[k] for k in id_list if k in embeddings_dict}, f)\n raise\n \n print(f\"[GPU {gpu_id}] āœ“ Completed {processed_count} sequences\")\n\n # Multi-GPU processing\n if num_gpus >= 2 and len(short_ids) > 1000: # Only split if enough sequences\n print(f\"\\n{'='*60}\")\n print(f\"PROCESSING SHORT SEQUENCES ON {num_gpus} GPUs (PARALLEL)\")\n print(f\"{'='*60}\")\n \n # Split short_ids across GPUs\n split_size = len(short_ids) // num_gpus\n gpu_splits = []\n for i in range(num_gpus):\n start = i * split_size\n end = start + split_size if i < num_gpus - 1 else len(short_ids)\n gpu_splits.append(short_ids[start:end])\n print(f\"GPU {i}: {len(gpu_splits[i])} sequences\")\n \n # Shared dictionary with thread lock\n from threading import Lock\n lock = Lock()\n \n # Create threads for each GPU\n threads = []\n for gpu_id in range(num_gpus):\n checkpoint_file_gpu = f'{CHECKPOINT_FILE}_gpu{gpu_id}.pkl'\n thread = Thread(\n target=process_on_gpu,\n args=(gpu_id, gpu_splits[gpu_id], train_embeddings, lock, checkpoint_file_gpu)\n )\n threads.append(thread)\n \n # Start all threads\n for thread in threads:\n thread.start()\n \n # Wait for completion\n for thread in threads:\n thread.join()\n \n print(f\"\\nāœ“ All GPUs completed short sequences\")\n \n # Cleanup GPU checkpoints\n for gpu_id in range(num_gpus):\n checkpoint_file_gpu = f'{CHECKPOINT_FILE}_gpu{gpu_id}.pkl'\n if os.path.exists(checkpoint_file_gpu):\n os.remove(checkpoint_file_gpu)\n \n else:\n # Single GPU processing\n print(f\"\\n{'='*60}\")\n print(f\"PROCESSING SHORT SEQUENCES ON SINGLE GPU\")\n print(f\"{'='*60}\")\n \n if len(short_ids) > 0:\n from threading import Lock\n lock = Lock()\n device_id = 0\n process_on_gpu(device_id, short_ids, train_embeddings, lock, f'{CHECKPOINT_FILE}_gpu0.pkl')\n\n # Process long sequences on GPUs (batch_size=1 for safety)\n if len(long_ids) > 0:\n print(f\"\\n{'='*60}\")\n print(f\"PROCESSING LONG SEQUENCES ON {num_gpus} GPU(S)\")\n print(f\"{'='*60}\")\n print(\"āš ļø Using batch_size=1 for long sequences to avoid OOM\")\n \n if num_gpus >= 2:\n # Split long sequences across GPUs\n split_size = len(long_ids) // num_gpus\n gpu_splits = []\n for i in range(num_gpus):\n start = i * split_size\n end = start + split_size if i < num_gpus - 1 else len(long_ids)\n gpu_splits.append(long_ids[start:end])\n print(f\"GPU {i}: {len(gpu_splits[i])} long sequences\")\n \n from threading import Lock\n lock = Lock()\n \n # Create threads for each GPU\n threads = []\n for gpu_id in range(num_gpus):\n checkpoint_file_gpu = f'{CHECKPOINT_FILE}_long_gpu{gpu_id}.pkl'\n thread = Thread(\n target=process_on_gpu,\n args=(gpu_id, gpu_splits[gpu_id], train_embeddings, lock, checkpoint_file_gpu)\n )\n threads.append(thread)\n \n # Start all threads\n for thread in threads:\n thread.start()\n \n # Wait for completion\n for thread in threads:\n thread.join()\n \n print(f\"\\nāœ“ All GPUs completed long sequences\")\n \n # Cleanup\n for gpu_id in range(num_gpus):\n checkpoint_file_gpu = f'{CHECKPOINT_FILE}_long_gpu{gpu_id}.pkl'\n if os.path.exists(checkpoint_file_gpu):\n os.remove(checkpoint_file_gpu)\n else:\n # Single GPU\n from threading import Lock\n lock = Lock()\n process_on_gpu(0, long_ids, train_embeddings, lock, f'{CHECKPOINT_FILE}_long_gpu0.pkl')\n\n # Final save\n print(f\"\\n{'='*60}\")\n print(f\"SAVING FINAL EMBEDDINGS\")\n print(f\"{'='*60}\")\n \n print(f\"Saving embeddings to {TRAIN_EMB_FILE}\")\n with open(TRAIN_EMB_FILE, 'wb') as f:\n pickle.dump(train_embeddings, f)\n \n if os.path.exists(CHECKPOINT_FILE):\n os.remove(CHECKPOINT_FILE)\n \n print(f\"āœ“ Extracted and saved {len(train_embeddings)} embeddings\")\n\nprint(f\"\\nāœ“ Training embeddings ready\")\nprint(f\" Embedding dimension: {list(train_embeddings.values())[0].shape[0]}\")\nprint(f\" Total proteins: {len(train_embeddings)}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Preparing training data...\")\n\n# Get proteins with both embeddings and annotations\ncommon_proteins = set(train_embeddings.keys()) & set(protein_terms.keys())\nprotein_ids = sorted(list(common_proteins))\n\nprint(f\"Proteins with both embeddings and annotations: {len(protein_ids)}\")\n\n# Create feature matrix X\nX = np.array([train_embeddings[pid] for pid in protein_ids])\n\n# Create multi-label target y\nmlb = MultiLabelBinarizer(classes=all_go_terms)\ny_labels = [protein_terms[pid] for pid in protein_ids]\ny = mlb.fit_transform(y_labels)\n\nprint(f\"\\nTraining data shape:\")\nprint(f\" X: {X.shape} (samples Ɨ embedding_dim)\")\nprint(f\" y: {y.shape} (samples Ɨ GO_terms)\")\nprint(f\" Average GO terms per protein: {y.sum(axis=1).mean():.2f}\")\nprint(f\" Sparsity: {(1 - y.mean())*100:.2f}% (mostly zeros)\")\n\nprint(\"\\nāœ“ Training data prepared\")\nprint(\"Hello World!!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Robust ESM embedding extraction with checkpoint resume and immediate short-save\n# Requires: pip install fair-esm\nimport os\nimport pickle\nimport gc\nimport time\nfrom tqdm.auto import tqdm\nimport torch\nimport esm\nfrom threading import Thread, Lock\nfrom queue import Queue\nimport glob\n\n# === User variables ===\nEMBEDDINGS_DIR = globals().get(\"EMBEDDINGS_DIR\", \"/kaggle/working/embeddings\")\nos.makedirs(EMBEDDINGS_DIR, exist_ok=True)\nTRAIN_EMB_FILE = os.path.join(EMBEDDINGS_DIR, 'train_embeddings.pkl') # final merged file\nTRAIN_SHORT_EMB_FILE = os.path.join(EMBEDDINGS_DIR, 'train_embeddings_short.pkl') # saved after short pass\nCHECKPOINT_PREFIX = os.path.join(EMBEDDINGS_DIR, 'train_embeddings_checkpoint') # checkpoint files will have suffixes\n\nMODEL_NAME = \"esm2_t30_150M_UR50D\"\nUSE_HALF = True\nBATCH_SIZE = 4\nCPU_BATCH_FOR_LONG = 1\nLONG_SEQ_THRESHOLD = 1000\nCHECKPOINT_EVERY = 500\nMAX_SEQ_LENGTH = 1022\n\n# === sanity: user must supply train_sequences dict (protein_id -> sequence) ===\nif 'train_sequences' not in globals():\n raise RuntimeError(\"Please define `train_sequences` (dict mapping protein_id -> sequence) before running this cell.\")\n\n# GPU check\nnum_gpus = torch.cuda.device_count()\nprint(f\"āœ“ Found {num_gpus} GPU(s)\")\n\nos.environ.setdefault(\"PYTORCH_CUDA_ALLOC_CONF\", \"expandable_segments:True\")\n\n# --- Helper functions for checkpointing & recovery ---\ndef save_pickle_atomic(obj, path):\n tmp = path + '.tmp'\n with open(tmp, 'wb') as f:\n pickle.dump(obj, f)\n os.replace(tmp, path)\n\ndef find_checkpoint_files(prefix=CHECKPOINT_PREFIX):\n # e.g. /kaggle/working/embeddings/train_embeddings_checkpoint_gpu0.pkl\n pattern = f\"{prefix}*gpu*.pkl\"\n return glob.glob(pattern)\n\ndef load_pickles(paths):\n merged = {}\n for p in paths:\n try:\n with open(p, 'rb') as f:\n data = pickle.load(f)\n if not isinstance(data, dict):\n print(f\"Warning: checkpoint {p} is not a dict - skipping\")\n continue\n merged.update(data)\n print(f\"Loaded {len(data)} entries from {p}\")\n except Exception as e:\n print(f\"Failed to load {p}: {e}\")\n return merged\n\ndef recover_embeddings():\n \"\"\"\n Attempt to reconstruct train_embeddings from:\n 1. final TRAIN_EMB_FILE\n 2. TRAIN_SHORT_EMB_FILE (short-only saved)\n 3. any per-gpu checkpoint files\n Return dict\n \"\"\"\n if os.path.exists(TRAIN_EMB_FILE):\n print(f\"Loading final embeddings from {TRAIN_EMB_FILE}\")\n with open(TRAIN_EMB_FILE, 'rb') as f:\n return pickle.load(f)\n\n recovered = {}\n # 1) short file\n if os.path.exists(TRAIN_SHORT_EMB_FILE):\n try:\n with open(TRAIN_SHORT_EMB_FILE, 'rb') as f:\n d = pickle.load(f)\n if isinstance(d, dict):\n recovered.update(d)\n print(f\"Recovered {len(d)} short embeddings from {TRAIN_SHORT_EMB_FILE}\")\n except Exception as e:\n print(f\"Could not load {TRAIN_SHORT_EMB_FILE}: {e}\")\n\n # 2) per-gpu checkpoint files\n ckpts = find_checkpoint_files()\n if ckpts:\n print(f\"Found checkpoint files: {ckpts}\")\n d = load_pickles(ckpts)\n recovered.update(d)\n\n # 3) final checkpoint (legacy name)\n legacy_ckpt = f\"{CHECKPOINT_PREFIX}.pkl\"\n if os.path.exists(legacy_ckpt):\n try:\n with open(legacy_ckpt, 'rb') as f:\n d = pickle.load(f)\n if isinstance(d, dict):\n recovered.update(d)\n print(f\"Recovered {len(d)} entries from legacy checkpoint {legacy_ckpt}\")\n except Exception as e:\n print(f\"Could not load {legacy_ckpt}: {e}\")\n\n print(f\"Total recovered embeddings: {len(recovered)}\")\n return recovered\n\n# Attempt to recover previously computed embeddings\ntrain_embeddings = recover_embeddings()\n\n# Build ID lists and pending work\nall_ids = list(train_sequences.keys())\npending_ids = [pid for pid in all_ids if pid not in train_embeddings]\ntotal_pending = len(pending_ids)\nprint(f\"Total sequences: {len(all_ids)}; Remaining to process: {total_pending}\")\n\nshort_ids = [pid for pid in pending_ids if len(train_sequences[pid]) < LONG_SEQ_THRESHOLD]\nlong_ids = [pid for pid in pending_ids if len(train_sequences[pid]) >= LONG_SEQ_THRESHOLD]\nprint(f\" Short (<{LONG_SEQ_THRESHOLD}): {len(short_ids)} Long (ā‰„{LONG_SEQ_THRESHOLD}): {len(long_ids)}\")\n\n# === processing function (thread safe) ===\ndef process_on_gpu(gpu_id, id_list, embeddings_dict, lock, checkpoint_file_gpu, is_short_phase=True):\n \"\"\"Process sequences on a specific GPU. Saves per-gpu checkpoint to checkpoint_file_gpu.\"\"\"\n device = torch.device(f'cuda:{gpu_id}') if torch.cuda.is_available() and gpu_id is not None else torch.device('cpu')\n print(f\"[GPU {gpu_id}] Starting with {len(id_list)} sequences; device={device}\")\n\n # Load model on this GPU\n model, alphabet = esm.pretrained.load_model_and_alphabet(MODEL_NAME)\n batch_converter = alphabet.get_batch_converter()\n model = model.to(device)\n model.eval()\n\n if USE_HALF and device.type == 'cuda':\n try:\n model.half()\n print(f\"[GPU {gpu_id}] Converted to float16\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Warning: half() failed: {e}\")\n\n idx = 0\n N = len(id_list)\n bs = BATCH_SIZE if is_short_phase else 1 # for long sequences, use batch=1 by default\n oom_count = 0\n processed_count = 0\n id_list.sort(key=lambda pid: len(train_sequences[pid]))\n\n while idx < N:\n batch_ids = id_list[idx: idx + bs]\n # clip to MAX_SEQ_LENGTH\n batch_data = [(pid, train_sequences[pid][:MAX_SEQ_LENGTH]) for pid in batch_ids]\n\n try:\n _, _, tokens = batch_converter(batch_data)\n tokens = tokens.to(device)\n\n if USE_HALF and device.type == 'cuda':\n context = torch.amp.autocast(device_type='cuda', dtype=torch.float16)\n else:\n class _null_ctx:\n def __enter__(self): return None\n def __exit__(self, exc_type, exc, tb): return False\n context = _null_ctx()\n\n with torch.no_grad():\n with context:\n results = model(tokens, repr_layers=[model.num_layers], return_contacts=False)\n\n rep_key = list(results[\"representations\"].keys())[-1]\n token_representations = results[\"representations\"][rep_key]\n # mean of token representations; shape (batch, dim)\n embeddings = token_representations.mean(dim=1).detach().cpu().float().numpy()\n\n with lock:\n for j, pid in enumerate(batch_ids):\n embeddings_dict[pid] = embeddings[j]\n processed_count += len(batch_ids)\n\n # Frequent checkpointing\n if processed_count % CHECKPOINT_EVERY == 0 or idx + bs >= N:\n # save only entries from this id_list to avoid partial overlap\n to_save = {k: embeddings_dict[k] for k in id_list if k in embeddings_dict}\n try:\n save_pickle_atomic(to_save, checkpoint_file_gpu)\n print(f\"[GPU {gpu_id}] Checkpoint saved ({len(to_save)} entries) -> {checkpoint_file_gpu}\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Failed to write checkpoint: {e}\")\n\n idx += bs\n oom_count = 0\n\n # free\n del tokens, results, token_representations, embeddings\n torch.cuda.empty_cache()\n gc.collect()\n\n except RuntimeError as e:\n msg = str(e).lower()\n if 'out of memory' in msg:\n oom_count += 1\n old_bs = bs\n bs = max(1, bs // 2)\n print(f\"[GPU {gpu_id}] OOM (batch {old_bs}). Reducing to {bs}\")\n\n if 'tokens' in locals():\n del tokens\n torch.cuda.empty_cache()\n gc.collect()\n\n if bs == 1 and oom_count >= 2 and device.type == 'cuda':\n print(f\"[GPU {gpu_id}] Persistent OOM. Switching to CPU for remaining.\")\n device = torch.device('cpu')\n model = model.cpu().float()\n torch.cuda.empty_cache()\n gc.collect()\n continue\n else:\n raise\n\n except KeyboardInterrupt:\n print(f\"[GPU {gpu_id}] Interrupted - saving checkpoint...\")\n with lock:\n to_save = {k: embeddings_dict[k] for k in id_list if k in embeddings_dict}\n save_pickle_atomic(to_save, checkpoint_file_gpu)\n raise\n\n print(f\"[GPU {gpu_id}] āœ“ Completed {processed_count} sequences\")\n # final save for this GPU\n with lock:\n to_save = {k: embeddings_dict[k] for k in id_list if k in embeddings_dict}\n try:\n save_pickle_atomic(to_save, checkpoint_file_gpu)\n print(f\"[GPU {gpu_id}] Final per-GPU checkpoint saved: {checkpoint_file_gpu} ({len(to_save)})\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Failed final save: {e}\")\n\n\n# === Short sequences processing (parallel if possible) ===\nif len(short_ids) > 0:\n print(f\"\\n{'='*60}\\nPROCESSING SHORT SEQUENCES\\n{'='*60}\")\n\n # If there are previously saved short embeddings, skip those IDs\n short_ids = [pid for pid in short_ids if pid not in train_embeddings]\n print(f\"Short sequences to process after recovery: {len(short_ids)}\")\n\n if len(short_ids) > 0:\n # Multi-GPU split conditions (keep simple)\n if num_gpus >= 2 and len(short_ids) > 1000:\n split_size = len(short_ids) // num_gpus\n gpu_splits = []\n for i in range(num_gpus):\n start = i * split_size\n end = start + split_size if i < num_gpus - 1 else len(short_ids)\n gpu_splits.append(short_ids[start:end])\n\n lock = Lock()\n threads = []\n for gpu_id in range(num_gpus):\n ckpt_path = f\"{CHECKPOINT_PREFIX}_gpu{gpu_id}.pkl\"\n t = Thread(target=process_on_gpu, args=(gpu_id, gpu_splits[gpu_id], train_embeddings, lock, ckpt_path, True))\n threads.append(t)\n\n for t in threads:\n t.start()\n for t in threads:\n t.join()\n\n # merge per-gpu checkpoints into train_embeddings (defensive)\n ckpts = [f\"{CHECKPOINT_PREFIX}_gpu{gpu_id}.pkl\" for gpu_id in range(num_gpus) if os.path.exists(f\"{CHECKPOINT_PREFIX}_gpu{gpu_id}.pkl\")]\n if ckpts:\n recovered_short = load_pickles(ckpts)\n train_embeddings.update(recovered_short)\n else:\n # single GPU or small dataset\n lock = Lock()\n ckpt_path = f\"{CHECKPOINT_PREFIX}_gpu0.pkl\"\n process_on_gpu(0, short_ids, train_embeddings, lock, ckpt_path, True)\n # load checkpoint to ensure consistent content\n if os.path.exists(ckpt_path):\n with open(ckpt_path, 'rb') as f:\n d = pickle.load(f)\n train_embeddings.update(d)\n\n # After short phase, persist a dedicated short-file to avoid losing this progress:\n try:\n short_save = {k: train_embeddings[k] for k in list(train_embeddings.keys()) if len(train_sequences[k]) < LONG_SEQ_THRESHOLD}\n save_pickle_atomic(short_save, TRAIN_SHORT_EMB_FILE)\n print(f\"āœ“ Saved short embeddings ({len(short_save)}) to {TRAIN_SHORT_EMB_FILE}\")\n except Exception as e:\n print(f\"Failed to save short embeddings file: {e}\")\n else:\n print(\"No short sequences to process (all recovered).\")\nelse:\n print(\"No short sequences present in dataset.\")\n\n# === Long sequences processing ===\nif len(long_ids) > 0:\n print(f\"\\n{'='*60}\\nPROCESSING LONG SEQUENCES\\n{'='*60}\")\n # Filter out ones recovered earlier:\n long_ids = [pid for pid in long_ids if pid not in train_embeddings]\n print(f\"Long sequences to process after recovery: {len(long_ids)}\")\n if len(long_ids) > 0:\n lock = Lock()\n if num_gpus >= 2:\n split_size = len(long_ids) // num_gpus\n gpu_splits = []\n for i in range(num_gpus):\n start = i * split_size\n end = start + split_size if i < num_gpus - 1 else len(long_ids)\n gpu_splits.append(long_ids[start:end])\n\n threads = []\n for gpu_id in range(num_gpus):\n ckpt_path = f\"{CHECKPOINT_PREFIX}_long_gpu{gpu_id}.pkl\"\n t = Thread(target=process_on_gpu, args=(gpu_id, gpu_splits[gpu_id], train_embeddings, lock, ckpt_path, False))\n threads.append(t)\n\n for t in threads: t.start()\n for t in threads: t.join()\n\n # merge per-gpu long checkpoints\n ckpts = [f\"{CHECKPOINT_PREFIX}_long_gpu{gpu_id}.pkl\" for gpu_id in range(num_gpus) if os.path.exists(f\"{CHECKPOINT_PREFIX}_long_gpu{gpu_id}.pkl\")]\n if ckpts:\n recovered_long = load_pickles(ckpts)\n train_embeddings.update(recovered_long)\n else:\n # single GPU CPU fallback\n ckpt_path = f\"{CHECKPOINT_PREFIX}_long_gpu0.pkl\"\n process_on_gpu(0, long_ids, train_embeddings, lock, ckpt_path, False)\n if os.path.exists(ckpt_path):\n with open(ckpt_path, 'rb') as f:\n d = pickle.load(f)\n train_embeddings.update(d)\n\n else:\n print(\"No long sequences to process (all recovered).\")\nelse:\n print(\"No long sequences present in dataset.\")\n\n# === Final save: merge short and long and write single file ===\nprint(f\"\\n{'='*60}\\nSAVING FINAL MERGED EMBEDDINGS\\n{'='*60}\")\ntry:\n # final reduce to only sequences present in all_ids (safety)\n final = {k: train_embeddings[k] for k in all_ids if k in train_embeddings}\n save_pickle_atomic(final, TRAIN_EMB_FILE)\n print(f\"āœ“ Final embeddings saved to {TRAIN_EMB_FILE} ({len(final)} entries)\")\n # optionally remove all per-gpu checkpoint files now that final saved\n for p in find_checkpoint_files():\n try:\n os.remove(p)\n except Exception:\n pass\n # also keep TRAIN_SHORT_EMB_FILE as additional backup (do not delete by default)\nexcept Exception as e:\n print(f\"Failed to save final embeddings: {e}\")\n\nprint(f\"\\nāœ“ Done. Total embeddings available: {len(final)}\")\nif len(final) > 0:\n print(f\" Embedding dimension: {list(final.values())[0].shape[0]}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Robust ESM embedding extraction with checkpoint resume and immediate short-save\n# Requires: pip install fair-esm\nimport os\nimport pickle\nimport gc\nimport time\nfrom tqdm.auto import tqdm\nimport torch\nimport esm\nfrom threading import Thread, Lock\nfrom queue import Queue\nimport glob\n\n# === User variables ===\nEMBEDDINGS_DIR = globals().get(\"EMBEDDINGS_DIR\", \"/kaggle/working/embeddings\")\nos.makedirs(EMBEDDINGS_DIR, exist_ok=True)\nTEST_EMB_FILE = os.path.join(EMBEDDINGS_DIR, 'test_embeddings.pkl') # final merged file\nTEST_SHORT_EMB_FILE = os.path.join(EMBEDDINGS_DIR, 'test_embeddings_short.pkl') # saved after short pass\nCHECKPOINT_PREFIX = os.path.join(EMBEDDINGS_DIR, 'test_embeddings_checkpoint') # checkpoint files will have suffixes\n\nMODEL_NAME = \"esm2_t30_150M_UR50D\"\nUSE_HALF = True\nBATCH_SIZE = 4\nCPU_BATCH_FOR_LONG = 1\nLONG_SEQ_THRESHOLD = 1000\nCHECKPOINT_EVERY = 500\nMAX_SEQ_LENGTH = 1022\n\n# === sanity: user must supply test_sequences dict (protein_id -> sequence) ===\nif 'test_sequences' not in globals():\n raise RuntimeError(\"Please define `test_sequences` (dict mapping protein_id -> sequence) before running this cell.\")\n\n# GPU check\nnum_gpus = torch.cuda.device_count()\nprint(f\"āœ“ Found {num_gpus} GPU(s)\")\n\nos.environ.setdefault(\"PYTORCH_CUDA_ALLOC_CONF\", \"expandable_segments:True\")\n\n# --- Helper functions for checkpointing & recovery ---\ndef save_pickle_atomic(obj, path):\n tmp = path + '.tmp'\n with open(tmp, 'wb') as f:\n pickle.dump(obj, f)\n os.replace(tmp, path)\n\ndef find_checkpoint_files(prefix=CHECKPOINT_PREFIX):\n # e.g. /kaggle/working/embeddings/test_embeddings_checkpoint_gpu0.pkl\n pattern = f\"{prefix}*gpu*.pkl\"\n return glob.glob(pattern)\n\ndef load_pickles(paths):\n merged = {}\n for p in paths:\n try:\n with open(p, 'rb') as f:\n data = pickle.load(f)\n if not isinstance(data, dict):\n print(f\"Warning: checkpoint {p} is not a dict - skipping\")\n continue\n merged.update(data)\n print(f\"Loaded {len(data)} entries from {p}\")\n except Exception as e:\n print(f\"Failed to load {p}: {e}\")\n return merged\n\ndef recover_embeddings():\n \"\"\"\n Attempt to reconstruct test_embeddings from:\n 1. final TEST_EMB_FILE\n 2. TEST_SHORT_EMB_FILE (short-only saved)\n 3. any per-gpu checkpoint files\n Return dict\n \"\"\"\n if os.path.exists(TEST_EMB_FILE):\n print(f\"Loading final embeddings from {TEST_EMB_FILE}\")\n with open(TEST_EMB_FILE, 'rb') as f:\n return pickle.load(f)\n\n recovered = {}\n # 1) short file\n if os.path.exists(TEST_SHORT_EMB_FILE):\n try:\n with open(TEST_SHORT_EMB_FILE, 'rb') as f:\n d = pickle.load(f)\n if isinstance(d, dict):\n recovered.update(d)\n print(f\"Recovered {len(d)} short embeddings from {TEST_SHORT_EMB_FILE}\")\n except Exception as e:\n print(f\"Could not load {TEST_SHORT_EMB_FILE}: {e}\")\n\n # 2) per-gpu checkpoint files\n ckpts = find_checkpoint_files()\n if ckpts:\n print(f\"Found checkpoint files: {ckpts}\")\n d = load_pickles(ckpts)\n recovered.update(d)\n\n # 3) final checkpoint (legacy name)\n legacy_ckpt = f\"{CHECKPOINT_PREFIX}.pkl\"\n if os.path.exists(legacy_ckpt):\n try:\n with open(legacy_ckpt, 'rb') as f:\n d = pickle.load(f)\n if isinstance(d, dict):\n recovered.update(d)\n print(f\"Recovered {len(d)} entries from legacy checkpoint {legacy_ckpt}\")\n except Exception as e:\n print(f\"Could not load {legacy_ckpt}: {e}\")\n\n print(f\"Total recovered embeddings: {len(recovered)}\")\n return recovered\n\n# Attempt to recover previously computed embeddings\ntest_embeddings = recover_embeddings()\n\n# Build ID lists and pending work\nall_ids = list(test_sequences.keys())\npending_ids = [pid for pid in all_ids if pid not in test_embeddings]\ntotal_pending = len(pending_ids)\nprint(f\"Total sequences: {len(all_ids)}; Remaining to process: {total_pending}\")\n\nshort_ids = [pid for pid in pending_ids if len(test_sequences[pid]) < LONG_SEQ_THRESHOLD]\nlong_ids = [pid for pid in pending_ids if len(test_sequences[pid]) >= LONG_SEQ_THRESHOLD]\nprint(f\" Short (<{LONG_SEQ_THRESHOLD}): {len(short_ids)} Long (ā‰„{LONG_SEQ_THRESHOLD}): {len(long_ids)}\")\n\n# === processing function (thread safe) ===\ndef process_on_gpu(gpu_id, id_list, embeddings_dict, lock, checkpoint_file_gpu, is_short_phase=True):\n \"\"\"Process sequences on a specific GPU. Saves per-gpu checkpoint to checkpoint_file_gpu.\"\"\"\n device = torch.device(f'cuda:{gpu_id}') if torch.cuda.is_available() and gpu_id is not None else torch.device('cpu')\n print(f\"[GPU {gpu_id}] Starting with {len(id_list)} sequences; device={device}\")\n\n # Load model on this GPU\n model, alphabet = esm.pretrained.load_model_and_alphabet(MODEL_NAME)\n batch_converter = alphabet.get_batch_converter()\n model = model.to(device)\n model.eval()\n\n if USE_HALF and device.type == 'cuda':\n try:\n model.half()\n print(f\"[GPU {gpu_id}] Converted to float16\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Warning: half() failed: {e}\")\n\n idx = 0\n N = len(id_list)\n bs = BATCH_SIZE if is_short_phase else 1 # for long sequences, use batch=1 by default\n oom_count = 0\n processed_count = 0\n id_list.sort(key=lambda pid: len(test_sequences[pid]))\n\n while idx < N:\n batch_ids = id_list[idx: idx + bs]\n # clip to MAX_SEQ_LENGTH\n batch_data = [(pid, test_sequences[pid][:MAX_SEQ_LENGTH]) for pid in batch_ids]\n\n try:\n _, _, tokens = batch_converter(batch_data)\n tokens = tokens.to(device)\n\n if USE_HALF and device.type == 'cuda':\n context = torch.amp.autocast(device_type='cuda', dtype=torch.float16)\n else:\n class _null_ctx:\n def __enter__(self): return None\n def __exit__(self, exc_type, exc, tb): return False\n context = _null_ctx()\n\n with torch.no_grad():\n with context:\n results = model(tokens, repr_layers=[model.num_layers], return_contacts=False)\n\n rep_key = list(results[\"representations\"].keys())[-1]\n token_representations = results[\"representations\"][rep_key]\n # mean of token representations; shape (batch, dim)\n embeddings = token_representations.mean(dim=1).detach().cpu().float().numpy()\n\n with lock:\n for j, pid in enumerate(batch_ids):\n embeddings_dict[pid] = embeddings[j]\n processed_count += len(batch_ids)\n\n # Frequent checkpointing\n if processed_count % CHECKPOINT_EVERY == 0 or idx + bs >= N:\n # save only entries from this id_list to avoid partial overlap\n to_save = {k: embeddings_dict[k] for k in id_list if k in embeddings_dict}\n try:\n save_pickle_atomic(to_save, checkpoint_file_gpu)\n print(f\"[GPU {gpu_id}] Checkpoint saved ({len(to_save)} entries) -> {checkpoint_file_gpu}\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Failed to write checkpoint: {e}\")\n\n idx += bs\n oom_count = 0\n\n # free\n del tokens, results, token_representations, embeddings\n torch.cuda.empty_cache()\n gc.collect()\n\n except RuntimeError as e:\n msg = str(e).lower()\n if 'out of memory' in msg:\n oom_count += 1\n old_bs = bs\n bs = max(1, bs // 2)\n print(f\"[GPU {gpu_id}] OOM (batch {old_bs}). Reducing to {bs}\")\n\n if 'tokens' in locals():\n del tokens\n torch.cuda.empty_cache()\n gc.collect()\n\n if bs == 1 and oom_count >= 2 and device.type == 'cuda':\n print(f\"[GPU {gpu_id}] Persistent OOM. Switching to CPU for remaining.\")\n device = torch.device('cpu')\n model = model.cpu().float()\n torch.cuda.empty_cache()\n gc.collect()\n continue\n else:\n raise\n\n except KeyboardInterrupt:\n print(f\"[GPU {gpu_id}] Interrupted - saving checkpoint...\")\n with lock:\n to_save = {k: embeddings_dict[k] for k in id_list if k in embeddings_dict}\n save_pickle_atomic(to_save, checkpoint_file_gpu)\n raise\n\n print(f\"[GPU {gpu_id}] āœ“ Completed {processed_count} sequences\")\n # final save for this GPU\n with lock:\n to_save = {k: embeddings_dict[k] for k in id_list if k in embeddings_dict}\n try:\n save_pickle_atomic(to_save, checkpoint_file_gpu)\n print(f\"[GPU {gpu_id}] Final per-GPU checkpoint saved: {checkpoint_file_gpu} ({len(to_save)})\")\n except Exception as e:\n print(f\"[GPU {gpu_id}] Failed final save: {e}\")\n\n\n# === Short sequences processing (parallel if possible) ===\nif len(short_ids) > 0:\n print(f\"\\n{'='*60}\\nPROCESSING SHORT SEQUENCES\\n{'='*60}\")\n\n # If there are previously saved short embeddings, skip those IDs\n short_ids = [pid for pid in short_ids if pid not in test_embeddings]\n print(f\"Short sequences to process after recovery: {len(short_ids)}\")\n\n if len(short_ids) > 0:\n # Multi-GPU split conditions (keep simple)\n if num_gpus >= 2 and len(short_ids) > 1000:\n split_size = len(short_ids) // num_gpus\n gpu_splits = []\n for i in range(num_gpus):\n start = i * split_size\n end = start + split_size if i < num_gpus - 1 else len(short_ids)\n gpu_splits.append(short_ids[start:end])\n\n lock = Lock()\n threads = []\n for gpu_id in range(num_gpus):\n ckpt_path = f\"{CHECKPOINT_PREFIX}_gpu{gpu_id}.pkl\"\n t = Thread(target=process_on_gpu, args=(gpu_id, gpu_splits[gpu_id], test_embeddings, lock, ckpt_path, True))\n threads.append(t)\n\n for t in threads:\n t.start()\n for t in threads:\n t.join()\n\n # merge per-gpu checkpoints into test_embeddings (defensive)\n ckpts = [f\"{CHECKPOINT_PREFIX}_gpu{gpu_id}.pkl\" for gpu_id in range(num_gpus) if os.path.exists(f\"{CHECKPOINT_PREFIX}_gpu{gpu_id}.pkl\")]\n if ckpts:\n recovered_short = load_pickles(ckpts)\n test_embeddings.update(recovered_short)\n else:\n # single GPU or small dataset\n lock = Lock()\n ckpt_path = f\"{CHECKPOINT_PREFIX}_gpu0.pkl\"\n process_on_gpu(0, short_ids, test_embeddings, lock, ckpt_path, True)\n # load checkpoint to ensure consistent content\n if os.path.exists(ckpt_path):\n with open(ckpt_path, 'rb') as f:\n d = pickle.load(f)\n test_embeddings.update(d)\n\n # After short phase, persist a dedicated short-file to avoid losing this progress:\n try:\n short_save = {k: test_embeddings[k] for k in list(test_embeddings.keys()) if len(test_sequences[k]) < LONG_SEQ_THRESHOLD}\n save_pickle_atomic(short_save, TEST_SHORT_EMB_FILE)\n print(f\"āœ“ Saved short embeddings ({len(short_save)}) to {TEST_SHORT_EMB_FILE}\")\n except Exception as e:\n print(f\"Failed to save short embeddings file: {e}\")\n else:\n print(\"No short sequences to process (all recovered).\")\nelse:\n print(\"No short sequences present in dataset.\")\n\n# === Long sequences processing ===\nif len(long_ids) > 0:\n print(f\"\\n{'='*60}\\nPROCESSING LONG SEQUENCES\\n{'='*60}\")\n # Filter out ones recovered earlier:\n long_ids = [pid for pid in long_ids if pid not in test_embeddings]\n print(f\"Long sequences to process after recovery: {len(long_ids)}\")\n if len(long_ids) > 0:\n lock = Lock()\n if num_gpus >= 2:\n split_size = len(long_ids) // num_gpus\n gpu_splits = []\n for i in range(num_gpus):\n start = i * split_size\n end = start + split_size if i < num_gpus - 1 else len(long_ids)\n gpu_splits.append(long_ids[start:end])\n\n threads = []\n for gpu_id in range(num_gpus):\n ckpt_path = f\"{CHECKPOINT_PREFIX}_long_gpu{gpu_id}.pkl\"\n t = Thread(target=process_on_gpu, args=(gpu_id, gpu_splits[gpu_id], test_embeddings, lock, ckpt_path, False))\n threads.append(t)\n\n for t in threads: t.start()\n for t in threads: t.join()\n\n # merge per-gpu long checkpoints\n ckpts = [f\"{CHECKPOINT_PREFIX}_long_gpu{gpu_id}.pkl\" for gpu_id in range(num_gpus) if os.path.exists(f\"{CHECKPOINT_PREFIX}_long_gpu{gpu_id}.pkl\")]\n if ckpts:\n recovered_long = load_pickles(ckpts)\n test_embeddings.update(recovered_long)\n else:\n # single GPU CPU fallback\n ckpt_path = f\"{CHECKPOINT_PREFIX}_long_gpu0.pkl\"\n process_on_gpu(0, long_ids, test_embeddings, lock, ckpt_path, False)\n if os.path.exists(ckpt_path):\n with open(ckpt_path, 'rb') as f:\n d = pickle.load(f)\n test_embeddings.update(d)\n\n else:\n print(\"No long sequences to process (all recovered).\")\nelse:\n print(\"No long sequences present in dataset.\")\n\n# === Final save: merge short and long and write single file ===\nprint(f\"\\n{'='*60}\\nSAVING FINAL MERGED EMBEDDINGS\\n{'='*60}\")\ntry:\n # final reduce to only sequences present in all_ids (safety)\n final = {k: test_embeddings[k] for k in all_ids if k in test_embeddings}\n save_pickle_atomic(final, TEST_EMB_FILE)\n print(f\"āœ“ Final embeddings saved to {TEST_EMB_FILE} ({len(final)} entries)\")\n # optionally remove all per-gpu checkpoint files now that final saved\n for p in find_checkpoint_files():\n try:\n os.remove(p)\n except Exception:\n pass\n # also keep TEST_SHORT_EMB_FILE as additional backup (do not delete by default)\nexcept Exception as e:\n print(f\"Failed to save final embeddings: {e}\")\n\nprint(f\"\\nāœ“ Done. Total embeddings available: {len(final)}\")\nif len(final) > 0:\n print(f\" Embedding dimension: {list(final.values())[0].shape[0]}\")\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MODEL_FILE = f'{MODEL_DIR}/logistic_model.pkl'\n\nif os.path.exists(MODEL_FILE):\n print(\"Loading trained model from cache...\")\n with open(MODEL_FILE, 'rb') as f:\n data = pickle.load(f)\n clf = data['model']\n \n mlb = data['mlb']\n print(\"āœ“ Model loaded from cache\")\nelse:\n print(\"Training multi-label classifier...\")\n print(\"āš ļø This will take 20-40 minutes\")\n print(\" (Go grab lunch! šŸ•)\")\n \n # Train One-vs-Rest Logistic Regression\n clf = MultiOutputClassifier(\n LogisticRegression(max_iter=1000, random_state=42, n_jobs=-1),\n n_jobs=-1\n )\n \n clf.fit(X, y)\n \n # Save model\n print(f\"\\nSaving model to {MODEL_FILE}\")\n with open(MODEL_FILE, 'wb') as f:\n pickle.dump({'model': clf, 'mlb': mlb}, f)\n \n print(\"āœ“ Model trained and saved\")\n\nprint(\"\\nāœ“ Classifier ready\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader, TensorDataset\nfrom sklearn.metrics import accuracy_score\n\nMODEL_FILE = f'{MODEL_DIR}/neural_model.pth'\n\nif os.path.exists(MODEL_FILE):\n print(\"Loading trained model from cache...\")\n checkpoint = torch.load(MODEL_FILE)\n \n # Rebuild model\n input_dim = X.shape[1]\n output_dim = y.shape[1]\n model = nn.Sequential(\n nn.Linear(input_dim, 2048),\n nn.ReLU(),\n nn.Dropout(0.3),\n nn.Linear(2048, 1024),\n nn.ReLU(),\n nn.Dropout(0.3),\n nn.Linear(1024, output_dim),\n nn.Sigmoid()\n )\n model.load_state_dict(checkpoint['model_state_dict'])\n mlb = checkpoint['mlb']\n print(\"āœ“ Model loaded from cache\")\nelse:\n print(\"Training neural network classifier on GPU...\")\n print(\"āš ļø This will take 10-20 minutes\")\n \n device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n print(f\"Using device: {device}\")\n \n # Prepare data\n X_train_t = torch.FloatTensor(X)\n y_train_t = torch.FloatTensor(y)\n \n dataset = TensorDataset(X_train_t, y_train_t)\n dataloader = DataLoader(dataset, batch_size=128, shuffle=True)\n \n # Build model\n input_dim = X.shape[1]\n output_dim = y.shape[1]\n \n model = nn.Sequential(\n nn.Linear(input_dim, 2048),\n nn.ReLU(),\n nn.Dropout(0.3),\n nn.Linear(2048, 1024),\n nn.ReLU(),\n nn.Dropout(0.3),\n nn.Linear(1024, output_dim),\n nn.Sigmoid()\n ).to(device)\n \n criterion = nn.BCELoss()\n optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n \n # Train\n epochs = 10\n model.train()\n \n for epoch in range(epochs):\n total_loss = 0\n for batch_X, batch_y in tqdm(dataloader, desc=f\"Epoch {epoch+1}/{epochs}\"):\n batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n \n optimizer.zero_grad()\n outputs = model(batch_X)\n loss = criterion(outputs, batch_y)\n loss.backward()\n optimizer.step()\n \n total_loss += loss.item()\n \n print(f\"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(dataloader):.4f}\")\n \n # Save model\n print(f\"\\nSaving model to {MODEL_FILE}\")\n torch.save({\n 'model_state_dict': model.state_dict(),\n 'mlb': mlb\n }, MODEL_FILE)\n \n print(\"āœ“ Model trained and saved\")\n\nprint(\"\\nāœ“ Neural network classifier ready\")\n\n# For predictions, use:\n# model.eval()\n# with torch.no_grad():\n# X_test_t = torch.FloatTensor(X_test).to(device)\n# predictions = model(X_test_t).cpu().numpy()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\n# Load checkpoint first to get dimensions\nMODEL_PATH = \"/kaggle/input/cafa6-files/neural_model.pth\"\ncheckpoint = torch.load(MODEL_PATH, map_location=\"cpu\", weights_only=False)\nstate_dict = checkpoint[\"model_state_dict\"]\nmlb = checkpoint[\"mlb\"]\n\n# Get dimensions from state_dict\ninput_dim = state_dict['0.weight'].shape[1] # Input features\noutput_dim = state_dict['6.weight'].shape[0] # Output GO terms\n\nprint(f\"Input dim: {input_dim}, Output dim: {output_dim}\")\n\n# Recreate model with EXACT same architecture (nn.Sequential)\nmodel = nn.Sequential(\n nn.Linear(input_dim, 2048),\n nn.ReLU(),\n nn.Dropout(0.3),\n nn.Linear(2048, 1024),\n nn.ReLU(),\n nn.Dropout(0.3),\n nn.Linear(1024, output_dim),\n nn.Sigmoid()\n)\n\n# Now load will work\nmodel.load_state_dict(state_dict)\nmodel.eval()\n\nprint(\"āœ… Model and label binarizer successfully loaded!\")\nprint(f\"Model has {sum(p.numel() for p in model.parameters())} parameters\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Generating predictions for test set...\")\n\n# Prepare test data\ntest_protein_ids = sorted(test_embeddings.keys())\nX_test = np.array([test_embeddings[pid] for pid in test_protein_ids])\nprint(f\"Predicting for {len(test_protein_ids)} test proteins...\")\n\n# Get probability predictions using neural network\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodel.eval()\nmodel = model.to(device)\n\n# Predict in batches to avoid memory issues\nbatch_size = 256\nall_predictions = []\n\nwith torch.no_grad():\n for i in tqdm(range(0, len(X_test), batch_size), desc=\"Predicting\"):\n batch_X = torch.FloatTensor(X_test[i:i+batch_size]).to(device)\n batch_preds = model(batch_X).cpu().numpy()\n all_predictions.append(batch_preds)\n\n# Concatenate all predictions\ny_pred_proba = np.vstack(all_predictions)\nprint(f\"Prediction shape: {y_pred_proba.shape}\")\n\n# Build predictions list\npredictions = []\ngo_terms = mlb.classes_\nmin_prob = 0.01\nmax_terms_per_protein = 1500\n\nprint(\"Processing predictions...\")\nfor i in tqdm(range(len(test_protein_ids))):\n protein_id = test_protein_ids[i]\n protein_probs = []\n \n for j, go_term in enumerate(go_terms):\n # Get probability (already in [0,1] from Sigmoid)\n prob = y_pred_proba[i, j]\n \n if prob > min_prob:\n protein_probs.append((go_term, prob))\n \n # Sort by probability and limit\n protein_probs.sort(key=lambda x: x[1], reverse=True)\n protein_probs = protein_probs[:max_terms_per_protein]\n \n for go_term, prob in protein_probs:\n predictions.append((protein_id, go_term, prob))\n\nprint(f\"\\nāœ“ Generated {len(predictions)} predictions\")\nprint(f\" Average predictions per protein: {len(predictions)/len(test_protein_ids):.1f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Creating submission file...\")\n\nSUBMISSION_FILE = 'submission.tsv'\n\nwith open(SUBMISSION_FILE, 'w') as f:\n for protein_id, go_term, prob in predictions:\n # Format to 3 significant figures\n prob_str = f\"{prob:.3g}\"\n f.write(f\"{protein_id}\\t{go_term}\\t{prob_str}\\n\")\n\nprint(f\"āœ“ Submission file created: {SUBMISSION_FILE}\")\nprint(f\" Total predictions: {len(predictions)}\")\nprint(f\"\\nšŸŽ‰ PIPELINE COMPLETE!\")\nprint(\"\\nNext steps:\")\nprint(\"1. Review submission.tsv\")\nprint(\"2. Submit to Kaggle competition\")\nprint(\"3. Check your leaderboard score\")\nprint(\"4. Iterate and improve!\")\n\n# Show sample predictions\nprint(\"\\nšŸ“‹ Sample predictions (first 10):\")\nfor i, (pid, term, prob) in enumerate(predictions[:10]):\n print(f\" {pid}\\t{term}\\t{prob:.3f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# CAFA-6: End-to-end pipeline (EDA ā†’ fast embeddings ā†’ baseline ā†’ submission)\nThis notebook is an end-to-end, **ready-to-run** Kaggle pipeline for CAFA-style protein function prediction.\nIt uses a **small ESM2 model** for fast embedding extraction (good for iteration), caches results,\nperforms cluster-aware validation, trains a One-vs-Rest logistic baseline, tunes thresholds with IA weights,\nand writes a submission.\n\nRun cells in order. Each code block is a single notebook cell.\n","metadata":{}},{"cell_type":"code","source":"# Run in a code cell. Might already be present on Kaggle. mmseqs2/diamond require separate system installs (instructions below).\n# Keep installs minimal and fast.\n!pip install --quiet fair-esm obonet joblib tqdm\n# optional (for faster training / advanced use) - lightgbm\n!pip install --quiet lightgbm\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Imports & config\nimport os, sys, time, math, glob, pickle, gc\nfrom collections import defaultdict, Counter\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport torch\nimport esm\nfrom sklearn.preprocessing import MultiLabelBinarizer, StandardScaler\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.cluster import MiniBatchKMeans\nfrom sklearn.metrics import f1_score\nimport joblib\n\n# CONFIG - set DATA_DIR to the Kaggle input folder for the competition dataset. /kaggle/input/cafa-6-protein-function-prediction\n# If you do not know the exact path, list /kaggle/input and edit accordingly.\nDATA_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\" # change if different\nWORK_DIR = \"/kaggle/working\"\nEMBED_DIR = os.path.join(WORK_DIR, \"embeddings\")\nos.makedirs(EMBED_DIR, exist_ok=True)\n\n# Embedding / model knobs (fast iteration)\nEMBED_MODEL = \"esm2_t6_8M_UR50D\" # small, fast model ā€” swap for final runs\nUSE_HALF = True\nBATCH_SIZE = 16\nMAX_SEQ_LEN = 1022\nCHUNK_SAVE_EVERY = 1000 # save per-chunk to prevent loss on preemption\n\n# Files (guessed standard competition filenames)\nTRAIN_FASTA = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\")\nTRAIN_TERMS = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\")\nTRAIN_TAXONOMY = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\")\nTEST_FASTA = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\nGO_OBO = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\")\nIA_FILE = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\")\nSAMPLE_SUB = os.path.join(DATA_DIR, \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\")\n\n# Embedding cache files\nEMB_NPZ = os.path.join(EMBED_DIR, \"esm_small_embeddings.npz\") # contains 'emb' and 'ids'\nEMB_MAP = os.path.join(EMBED_DIR, \"id_to_idx.pkl\")\n\nRANDOM_SEED = 42\n\n# quick check for existence\nfor f in [TRAIN_FASTA, TRAIN_TERMS, TEST_FASTA, GO_OBO, IA_FILE]:\n if not os.path.exists(f):\n print(\"Warning: expected file not found:\", f)\nprint(\"Working dir:\", WORK_DIR, \"Embedding cache dir:\", EMBED_DIR)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Utility helpers\n\ndef read_fasta_to_dict(path):\n \"\"\"Return dict: accession -> sequence. Uses header first token (supports UniProt style).\"\"\"\n seqs = {}\n if not os.path.exists(path):\n raise FileNotFoundError(path)\n with open(path) as fh:\n cur_id = None\n cur_seq = []\n for line in fh:\n line = line.strip()\n if not line:\n continue\n if line.startswith(\">\"):\n if cur_id:\n seqs[cur_id] = \"\".join(cur_seq)\n header = line[1:].split()\n # pick UniProt accession-style last token after '|' if present\n token = header[0]\n if \"|\" in token:\n cur_id = token.split(\"|\")[-1]\n else:\n cur_id = token\n cur_seq = []\n else:\n cur_seq.append(line)\n if cur_id:\n seqs[cur_id] = \"\".join(cur_seq)\n return seqs\n\ndef save_pickle_atomic(obj, path):\n tmp = path + \".tmp\"\n with open(tmp, \"wb\") as fh:\n pickle.dump(obj, fh)\n os.replace(tmp, path)\n\ndef load_terms(tsv_path):\n df = pd.read_csv(tsv_path, sep=\"\\t\", header=None, names=[\"prot\",\"go\",\"subont\"])\n terms = defaultdict(list)\n for _, r in df.iterrows():\n terms[r.prot].append(r.go)\n return terms, df\n\ndef load_ia(ia_path):\n df = pd.read_csv(ia_path, sep=\"\\t\", header=None, names=[\"go\",\"ia\"])\n return {r.go: float(r.ia) for _, r in df.iterrows()}\n\ndef load_go_parents(obo_path):\n parents = defaultdict(list)\n if not os.path.exists(obo_path):\n print(\"GO OBO not found:\", obo_path)\n return parents\n current = {}\n with open(obo_path) as fh:\n for line in fh:\n line = line.strip()\n if line == \"[Term]\":\n current = {}\n elif line == \"\" and current:\n if 'id' in current:\n goid = current['id']\n for p in current.get('is_a', []):\n parents[goid].append(p.split()[0])\n current = {}\n elif line.startswith(\"id:\"):\n current['id'] = line.split()[1]\n elif line.startswith(\"is_a:\"):\n current.setdefault('is_a', []).append(line.split()[1])\n return parents\n\ndef propagate_scores(scores, parents_map):\n \"\"\"Recursively propagate child scores to parents via max(child_score).\"\"\"\n memo = {}\n def get_score(t):\n if t in memo:\n return memo[t]\n s = scores.get(t, 0.0)\n for p in parents_map.get(t, []):\n s = max(s, get_score(p))\n memo[t] = s\n return s\n for t in list(scores.keys()):\n get_score(t)\n return memo\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load sequences and labels (safe wrappers)\nprint(\"Loading data... (this may take a few seconds)\")\ntrain_seqs = read_fasta_to_dict(TRAIN_FASTA) if os.path.exists(TRAIN_FASTA) else {}\ntest_seqs = read_fasta_to_dict(TEST_FASTA) if os.path.exists(TEST_FASTA) else {}\ntrain_terms, train_terms_df = (load_terms(TRAIN_TERMS) if os.path.exists(TRAIN_TERMS) else (defaultdict(list), None))\nia_map = load_ia(IA_FILE) if os.path.exists(IA_FILE) else {}\nparents_map = load_go_parents(GO_OBO) if os.path.exists(GO_OBO) else {}\n\nprint(\"train seqs:\", len(train_seqs), \"test superset seqs:\", len(test_seqs))\nprint(\"train annotations (unique proteins):\", len(train_terms))\nprint(\"IA weights loaded:\", len(ia_map), \"GO parents loaded:\", len(parents_map))\n\n# Quick EDA - sequence lengths and top GO terms\nall_lengths = np.array([len(s) for s in list(train_seqs.values()) + list(test_seqs.values())]) if (train_seqs or test_seqs) else np.array([0])\nprint(\"Seq length: min/median/mean/max:\", all_lengths.min(), np.median(all_lengths), all_lengths.mean(), all_lengths.max())\n\nif train_terms_df is not None:\n print(\"Top 15 GO terms:\")\n print(train_terms_df['go'].value_counts().head(15))\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}