{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31192,"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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install obonet Bio\n!pip install biopython pandas scipy numpy transformers ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T04:04:06.751366Z","iopub.execute_input":"2025-12-17T04:04:06.751597Z","iopub.status.idle":"2025-12-17T04:04:16.678693Z","shell.execute_reply.started":"2025-12-17T04:04:06.751573Z","shell.execute_reply":"2025-12-17T04:04:16.677960Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import os\n# import torch\n# import torch.nn as nn\n# import torch.optim as optim\n# from torch.utils.data import Dataset, DataLoader, random_split\n# from transformers import AutoTokenizer, EsmModel\n# from Bio import SeqIO\n# import obonet\n# import networkx as nx\n# import pandas as pd\n# import numpy as np\n# from tqdm import tqdm\n# from collections import Counter\n# import gc\n\n# # ==========================================\n# # CONFIGURATION\n# # ==========================================\n# class Config:\n#     # Kaggle Standard Paths\n#     BASE_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n    \n#     # Note: Kaggle datasets sometimes have slightly different internal structures.\n#     # We check if 'Train' exists, otherwise files might be in the root.\n#     if os.path.exists(os.path.join(BASE_DIR, 'Train')):\n#         TRAIN_FASTA = os.path.join(BASE_DIR, 'Train', 'train_sequences.fasta')\n#         TRAIN_TERMS = os.path.join(BASE_DIR, 'Train', 'train_terms.tsv')\n#         OBO_FILE    = os.path.join(BASE_DIR, 'Train', 'go-basic.obo')\n#     else:\n#         # Fallback if files are in root\n#         TRAIN_FASTA = os.path.join(BASE_DIR, 'train_sequences.fasta')\n#         TRAIN_TERMS = os.path.join(BASE_DIR, 'train_terms.tsv')\n#         OBO_FILE    = os.path.join(BASE_DIR, 'go-basic.obo')\n\n#     if os.path.exists(os.path.join(BASE_DIR, 'Test')):\n#         TEST_FASTA  = os.path.join(BASE_DIR, 'Test', 'testsuperset.fasta')\n#     else:\n#         TEST_FASTA  = os.path.join(BASE_DIR, 'testsuperset.fasta')\n    \n#     # Model Parameters\n#     # Use \"facebook/esm2_t6_8M_UR50D\" for debugging (Fast)\n#     # Use \"facebook/esm2_t33_650M_UR50D\" for better results (Slow)\n#     MODEL_NAME = \"facebook/esm2_t6_8M_UR50D\"\n#     MAX_SEQ_LEN = 1022  \n#     EMBED_DIM = 320     # Match this to your model (320 for 8M, 1280 for 650M)\n#     NUM_LABELS = 1500   # Competition limit is 1500 terms per protein\n#     BATCH_SIZE = 16     # Conservative batch size for GPU\n#     LEARNING_RATE = 1e-3\n#     EPOCHS = 25\n    \n#     # Toggle this: True = Fast run on 500 samples. False = Full run.\n#     DEBUG_MODE = True \n\n# # ==========================================\n# # 1. ONTOLOGY & LABEL PROCESSING\n# # ==========================================\n# class LabelManager:\n#     def __init__(self, obo_path, terms_path, num_labels=1500):\n#         print(f\">> Loading Ontology from {obo_path}...\")\n#         self.graph = obonet.read_obo(obo_path)\n        \n#         # Precompute ancestors for fast lookups (Graph Traversal)\n#         self.term_ancestors = {\n#             term: set(nx.ancestors(self.graph, term)) \n#             for term in tqdm(self.graph.nodes, desc=\"Indexing Ancestors\")\n#         }\n        \n#         print(f\">> Loading Terms from {terms_path}...\")\n#         self.terms_df = pd.read_csv(terms_path, sep=\"\\t\")\n        \n#         # FIX: Standardize column names\n#         # The file typically has 'EntryID', 'term', 'aspect' or similar\n#         print(f\"   Original Columns: {self.terms_df.columns.tolist()}\")\n        \n#         # Identify the protein ID column (usually the first one or named EntryID/ProteinID)\n#         # We rename it to 'protein_id' for consistency\n#         possible_id_cols = ['EntryID', 'ProteinID', 'UniprotID', 'protein_id']\n#         found_col = None\n#         for col in possible_id_cols:\n#             if col in self.terms_df.columns:\n#                 found_col = col\n#                 break\n        \n#         if found_col:\n#             self.terms_df.rename(columns={found_col: 'protein_id'}, inplace=True)\n#             print(f\"   Renamed '{found_col}' to 'protein_id'\")\n#         else:\n#             # Fallback: Assume first column is ID\n#             first_col = self.terms_df.columns[0]\n#             self.terms_df.rename(columns={first_col: 'protein_id'}, inplace=True)\n#             print(f\"   Renamed first column '{first_col}' to 'protein_id'\")\n            \n#         # Select Top-K most frequent terms to create our classification targets\n#         term_counts = Counter(self.terms_df['term'])\n#         self.selected_terms = [t[0] for t in term_counts.most_common(num_labels)]\n#         self.term_to_idx = {t: i for i, t in enumerate(self.selected_terms)}\n#         self.idx_to_term = {i: t for t, i in self.term_to_idx.items()}\n#         print(f\"Selected {len(self.selected_terms)} most frequent GO terms.\")\n\n#     def get_labels_dict(self):\n#         \"\"\"\n#         Returns {protein_id: binary_vector}\n#         Performs 'Label Propagation': If a protein has a specific term, \n#         it also implicitly has all parent terms.\n#         \"\"\"\n#         # Now 'protein_id' is guaranteed to exist\n#         protein_groups = self.terms_df.groupby('protein_id')['term'].apply(set)\n#         labels_dict = {}\n        \n#         print(\">> Propagating Labels (Hierarchy Enforcement)...\")\n        \n#         # In DEBUG_MODE, we don't want to process the whole label file, \n#         # but we need to ensure we keep labels for the proteins we actually load later.\n#         # So we process all labels here (it's fast enough) to avoid mismatches.\n#         all_ids = list(protein_groups.index)\n            \n#         for pid in tqdm(all_ids):\n#             raw_terms = protein_groups[pid]\n#             propagated_terms = set(raw_terms)\n            \n#             # Add all parents (The Rule: Child implies Parent)\n#             for term in raw_terms:\n#                 if term in self.term_ancestors:\n#                     propagated_terms.update(self.term_ancestors[term])\n            \n#             # Create Multi-Hot Binary Vector\n#             vec = np.zeros(Config.NUM_LABELS, dtype=np.float32)\n#             for term in propagated_terms:\n#                 if term in self.term_to_idx:\n#                     vec[self.term_to_idx[term]] = 1.0\n#             labels_dict[pid] = vec\n            \n#         return labels_dict\n\n# # ==========================================\n# # 2. FEATURE EXTRACTION (ESM-2)\n# # ==========================================\n# def clean_id(header_id):\n#     \"\"\"\n#     Fixes the Critical Bug:\n#     Input: \"sp|P9WHI7|RECN_MYCT\" or \"tr|...|...\"\n#     Output: \"P9WHI7\"\n#     \"\"\"\n#     if \"|\" in header_id:\n#         parts = header_id.split('|')\n#         # Standard SwissProt format is db|UniqueID|EntryName\n#         if len(parts) >= 2:\n#             return parts[1]\n#     return header_id\n\n# def extract_embeddings(fasta_path, model_name, limit=None):\n#     \"\"\"\n#     Reads FASTA -> ESM-2 -> Embeddings\n#     \"\"\"\n#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     print(f\">> Extracting Features using {model_name} on {device}...\")\n    \n#     tokenizer = AutoTokenizer.from_pretrained(model_name)\n#     model = EsmModel.from_pretrained(model_name).to(device)\n#     model.eval()\n    \n#     seqs = []\n#     ids = []\n    \n#     print(f\"Reading FASTA: {fasta_path}\")\n#     for i, record in enumerate(SeqIO.parse(fasta_path, \"fasta\")):\n#         if limit and i >= limit: break\n        \n#         # CRITICAL FIX: Clean the ID to match the TSV file\n#         clean_pid = clean_id(str(record.id))\n#         ids.append(clean_pid)\n        \n#         # Truncate sequence to model's limit\n#         seqs.append(str(record.seq)[:Config.MAX_SEQ_LEN]) \n        \n#     # Process in Batches\n#     embeddings = []\n#     # Smaller batch size for inference to prevent OOM\n#     batch_size = Config.BATCH_SIZE \n    \n#     with torch.no_grad():\n#         for i in tqdm(range(0, len(seqs), batch_size), desc=\"Embedding Batches\"):\n#             batch_seqs = seqs[i : i + batch_size]\n            \n#             # Tokenize\n#             inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True)\n#             inputs = {k: v.to(device) for k, v in inputs.items()}\n            \n#             # Forward Pass\n#             outputs = model(**inputs)\n            \n#             # Take the <CLS> token (position 0) as the sequence embedding\n#             batch_emb = outputs.last_hidden_state[:, 0, :].cpu().numpy()\n#             embeddings.append(batch_emb)\n            \n#     if len(embeddings) == 0:\n#         return [], np.array([])\n        \n#     return ids, np.vstack(embeddings)\n\n# # ==========================================\n# # 3. PYTORCH DATASET & MODEL\n# # ==========================================\n# class ProteinDataset(Dataset):\n#     def __init__(self, embeddings, labels=None):\n#         self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n#         self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n        \n#     def __len__(self):\n#         return len(self.embeddings)\n    \n#     def __getitem__(self, idx):\n#         if self.labels is not None:\n#             return self.embeddings[idx], self.labels[idx]\n#         return self.embeddings[idx]\n\n# class MultiLabelClassifier(nn.Module):\n#     def __init__(self, input_dim, num_classes):\n#         super().__init__()\n#         self.layer1 = nn.Linear(input_dim, 512)\n#         self.bn1 = nn.BatchNorm1d(512)\n#         self.relu = nn.ReLU()\n#         self.dropout = nn.Dropout(0.3)\n        \n#         self.layer2 = nn.Linear(512, 256) # Added an extra layer for depth\n#         self.bn2 = nn.BatchNorm1d(256)\n        \n#         self.output = nn.Linear(256, num_classes)\n        \n#     def forward(self, x):\n#         x = self.layer1(x)\n#         x = self.bn1(x)\n#         x = self.relu(x)\n#         x = self.dropout(x)\n        \n#         x = self.layer2(x)\n#         x = self.bn2(x)\n#         x = self.relu(x)\n#         x = self.dropout(x)\n        \n#         x = self.output(x)\n#         return x\n\n# # ==========================================\n# # 4. MAIN EXECUTION PIPELINE\n# # ==========================================\n# def main():\n#     print(\"=== STARTING CAFA 6 PIPELINE ===\")\n#     print(f\"Debug Mode: {Config.DEBUG_MODE}\")\n    \n#     # --- Step A: Prepare Labels ---\n#     label_mgr = LabelManager(Config.OBO_FILE, Config.TRAIN_TERMS, Config.NUM_LABELS)\n#     labels_map = label_mgr.get_labels_dict()\n    \n#     # --- Step B: Prepare Training Data ---\n#     train_limit = 500 if Config.DEBUG_MODE else None\n    \n#     # 1. Get X (Embeddings)\n#     train_ids, train_embeds = extract_embeddings(Config.TRAIN_FASTA, Config.MODEL_NAME, limit=train_limit)\n    \n#     # 2. Match X with Y (Labels)\n#     X_data = []\n#     y_data = []\n#     valid_ids = 0\n    \n#     print(\">> Aligning Embeddings with Labels...\")\n#     for pid, emb in zip(train_ids, train_embeds):\n#         if pid in labels_map:\n#             X_data.append(emb)\n#             y_data.append(labels_map[pid])\n#             valid_ids += 1\n            \n#     print(f\"Aligned {valid_ids} proteins out of {len(train_ids)} sequences.\")\n    \n#     if valid_ids == 0:\n#         print(\"ERROR: No IDs matched! Check FASTA header vs TSV protein_id format.\")\n#         return\n\n#     X_data = np.array(X_data)\n#     y_data = np.array(y_data)\n    \n#     # 3. Create DataLoaders\n#     dataset = ProteinDataset(X_data, y_data)\n#     train_size = int(0.8 * len(dataset))\n#     val_size = len(dataset) - train_size\n#     train_ds, val_ds = random_split(dataset, [train_size, val_size])\n    \n#     train_loader = DataLoader(train_ds, batch_size=Config.BATCH_SIZE, shuffle=True)\n#     val_loader = DataLoader(val_ds, batch_size=Config.BATCH_SIZE)\n    \n#     # --- Step C: Training the Model ---\n#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     model = MultiLabelClassifier(input_dim=Config.EMBED_DIM, num_classes=Config.NUM_LABELS).to(device)\n    \n#     # BCEWithLogitsLoss is standard for Multi-Label Classification\n#     criterion = nn.BCEWithLogitsLoss()\n#     optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)\n    \n#     print(f\"\\n>> Training for {Config.EPOCHS} epochs on {device}...\")\n    \n#     for epoch in range(Config.EPOCHS):\n#         model.train()\n#         train_loss = 0\n#         for inputs, targets in train_loader:\n#             inputs, targets = inputs.to(device), targets.to(device)\n#             optimizer.zero_grad()\n#             outputs = model(inputs)\n#             loss = criterion(outputs, targets)\n#             loss.backward()\n#             optimizer.step()\n#             train_loss += loss.item()\n            \n#         # Validation\n#         model.eval()\n#         val_loss = 0\n#         with torch.no_grad():\n#             for inputs, targets in val_loader:\n#                 inputs, targets = inputs.to(device), targets.to(device)\n#                 outputs = model(inputs)\n#                 loss = criterion(outputs, targets)\n#                 val_loss += loss.item()\n                \n#         print(f\"Epoch {epoch+1}/{Config.EPOCHS} | Train Loss: {train_loss/len(train_loader):.4f} | Val Loss: {val_loss/len(val_loader):.4f}\")\n\n#     torch.save(model.state_dict(), \"cafa6_model.pth\")\n#     print(\">> Model Saved!\")\n\n#     # --- Step D: Inference on Test Set ---\n#     print(\"\\n>> Processing Test Data...\")\n    \n#     test_limit = 200 if Config.DEBUG_MODE else None\n    \n#     # Note: Test set IDs also need cleaning\n#     test_ids, test_embeds = extract_embeddings(Config.TEST_FASTA, Config.MODEL_NAME, limit=test_limit)\n    \n#     if len(test_ids) == 0:\n#         print(\"No test data found/processed.\")\n#         return\n\n#     test_dataset = ProteinDataset(test_embeds, labels=None)\n#     test_loader = DataLoader(test_dataset, batch_size=Config.BATCH_SIZE, shuffle=False)\n    \n#     predictions = []\n#     model.eval()\n#     with torch.no_grad():\n#         for inputs in test_loader:\n#             inputs = inputs.to(device)\n#             outputs = model(inputs)\n#             # Apply Sigmoid to get 0..1 probabilities\n#             probs = torch.sigmoid(outputs).cpu().numpy()\n#             predictions.append(probs)\n            \n#     predictions = np.vstack(predictions)\n    \n#     # --- Step E: Submission File Generation ---\n#     print(\">> Generating submission.tsv...\")\n#     with open('submission.tsv', 'w') as f:\n#         for i, pid in enumerate(tqdm(test_ids)):\n#             probs = predictions[i]\n            \n#             # Select top predictions to reduce file size\n#             # Get indices of top 50 scores\n#             top_indices = np.argsort(probs)[-50:][::-1]\n            \n#             for idx in top_indices:\n#                 score = probs[idx]\n#                 if score > 0.01: # Confidence threshold\n#                     term_id = label_mgr.idx_to_term[idx]\n#                     f.write(f\"{pid}\\t{term_id}\\t{score:.3f}\\n\")\n                    \n#     print(\"=== PIPELINE COMPLETE ===\")\n\n# if __name__ == \"__main__\":\n#     main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T17:02:49.587789Z","iopub.execute_input":"2025-12-14T17:02:49.588727Z","iopub.status.idle":"2025-12-14T17:03:22.773439Z","shell.execute_reply.started":"2025-12-14T17:02:49.588698Z","shell.execute_reply":"2025-12-14T17:03:22.772756Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# CAFA 6 Single model Pipeline","metadata":{}},{"cell_type":"code","source":"# import os\n# import torch\n# import torch.nn as nn\n# import torch.optim as optim\n# from torch.utils.data import Dataset, DataLoader, random_split\n# from transformers import AutoTokenizer, EsmModel\n# from Bio import SeqIO\n# import obonet\n# import networkx as nx\n# import pandas as pd\n# import numpy as np\n# from tqdm import tqdm\n# from collections import Counter\n# import gc\n\n# # ==========================================\n# # CONFIGURATION\n# # ==========================================\n# class Config:\n#     # Paths\n#     BASE_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n#     if os.path.exists(os.path.join(BASE_DIR, 'Train')):\n#         TRAIN_FASTA = os.path.join(BASE_DIR, 'Train', 'train_sequences.fasta')\n#         TRAIN_TERMS = os.path.join(BASE_DIR, 'Train', 'train_terms.tsv')\n#         OBO_FILE    = os.path.join(BASE_DIR, 'Train', 'go-basic.obo')\n#     else:\n#         TRAIN_FASTA = os.path.join(BASE_DIR, 'train_sequences.fasta')\n#         TRAIN_TERMS = os.path.join(BASE_DIR, 'train_terms.tsv')\n#         OBO_FILE    = os.path.join(BASE_DIR, 'go-basic.obo')\n\n#     if os.path.exists(os.path.join(BASE_DIR, 'Test')):\n#         TEST_FASTA  = os.path.join(BASE_DIR, 'Test', 'testsuperset.fasta')\n#     else:\n#         TEST_FASTA  = os.path.join(BASE_DIR, 'testsuperset.fasta')\n    \n#     # --- UPGRADED MODEL ---\n#     # 35M parameters (better than 8M, faster than 650M)\n#     MODEL_NAME = \"facebook/esm2_t12_35M_UR50D\" \n#     EMBED_DIM = 480     # 480 for 35M, 320 for 8M, 640 for 150M, 1280 for 650M\n    \n#     MAX_SEQ_LEN = 1022  \n#     NUM_LABELS = 2000   # Increased to top 2000 terms\n#     BATCH_SIZE = 32     \n#     LEARNING_RATE = 5e-4 # Slightly lower starting rate for stability\n#     EPOCHS = 15         # 15 is usually enough with a good scheduler\n    \n#     # CRITICAL: Set to False for actual submission\n#     DEBUG_MODE = False \n\n# # ==========================================\n# # 1. ONTOLOGY & LABEL PROCESSING\n# # ==========================================\n# class LabelManager:\n#     def __init__(self, obo_path, terms_path, num_labels=1500):\n#         print(f\">> Loading Ontology from {obo_path}...\")\n#         self.graph = obonet.read_obo(obo_path)\n#         self.term_ancestors = {\n#             term: set(nx.ancestors(self.graph, term)) \n#             for term in tqdm(self.graph.nodes, desc=\"Indexing Ancestors\")\n#         }\n        \n#         print(f\">> Loading Terms from {terms_path}...\")\n#         self.terms_df = pd.read_csv(terms_path, sep=\"\\t\")\n        \n#         # ID Standardization\n#         possible_id_cols = ['EntryID', 'ProteinID', 'UniprotID', 'protein_id']\n#         found_col = None\n#         for col in possible_id_cols:\n#             if col in self.terms_df.columns:\n#                 found_col = col\n#                 break\n#         if found_col:\n#             self.terms_df.rename(columns={found_col: 'protein_id'}, inplace=True)\n#         else:\n#             self.terms_df.rename(columns={self.terms_df.columns[0]: 'protein_id'}, inplace=True)\n            \n#         term_counts = Counter(self.terms_df['term'])\n#         self.selected_terms = [t[0] for t in term_counts.most_common(num_labels)]\n#         self.term_to_idx = {t: i for i, t in enumerate(self.selected_terms)}\n#         self.idx_to_term = {i: t for t, i in self.term_to_idx.items()}\n#         print(f\"Selected {len(self.selected_terms)} most frequent GO terms.\")\n\n#     def get_labels_dict(self):\n#         protein_groups = self.terms_df.groupby('protein_id')['term'].apply(set)\n#         labels_dict = {}\n#         all_ids = list(protein_groups.index)\n        \n#         if Config.DEBUG_MODE: all_ids = all_ids[:500]\n            \n#         print(\">> Propagating Labels...\")\n#         for pid in tqdm(all_ids):\n#             raw_terms = protein_groups[pid]\n#             propagated_terms = set(raw_terms)\n#             for term in raw_terms:\n#                 if term in self.term_ancestors:\n#                     propagated_terms.update(self.term_ancestors[term])\n            \n#             vec = np.zeros(Config.NUM_LABELS, dtype=np.float32)\n#             for term in propagated_terms:\n#                 if term in self.term_to_idx:\n#                     vec[self.term_to_idx[term]] = 1.0\n#             labels_dict[pid] = vec\n#         return labels_dict\n\n# # ==========================================\n# # 2. FEATURE EXTRACTION\n# # ==========================================\n# def clean_id(header_id):\n#     if \"|\" in header_id:\n#         parts = header_id.split('|')\n#         if len(parts) >= 2: return parts[1]\n#     return header_id\n\n# def extract_embeddings(fasta_path, model_name, limit=None):\n#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     print(f\">> Extracting Features using {model_name} on {device}...\")\n    \n#     tokenizer = AutoTokenizer.from_pretrained(model_name)\n#     model = EsmModel.from_pretrained(model_name).to(device)\n#     model.eval()\n    \n#     seqs, ids = [], []\n#     for i, record in enumerate(SeqIO.parse(fasta_path, \"fasta\")):\n#         if limit and i >= limit: break\n#         ids.append(clean_id(str(record.id)))\n#         seqs.append(str(record.seq)[:Config.MAX_SEQ_LEN]) \n        \n#     embeddings = []\n#     # Smaller batch size for ESM inference to avoid OOM\n#     # 35M model is larger, so we reduce batch size slightly for safety\n#     batch_size = 8 \n    \n#     with torch.no_grad():\n#         for i in tqdm(range(0, len(seqs), batch_size), desc=\"Embed Batches\"):\n#             batch_seqs = seqs[i : i + batch_size]\n#             inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True)\n#             inputs = {k: v.to(device) for k, v in inputs.items()}\n#             outputs = model(**inputs)\n#             batch_emb = outputs.last_hidden_state[:, 0, :].cpu().numpy()\n#             embeddings.append(batch_emb)\n            \n#     if not embeddings: return [], np.array([])\n#     return ids, np.vstack(embeddings)\n\n# # ==========================================\n# # 3. ADVANCED MODEL & DATASET\n# # ==========================================\n# class ProteinDataset(Dataset):\n#     def __init__(self, embeddings, labels=None):\n#         self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n#         self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n#     def __len__(self): return len(self.embeddings)\n#     def __getitem__(self, idx):\n#         if self.labels is not None: return self.embeddings[idx], self.labels[idx]\n#         return self.embeddings[idx]\n\n# # IMPROVEMENT: Residual Block for deeper learning\n# class ResidualBlock(nn.Module):\n#     def __init__(self, dim, dropout=0.3):\n#         super().__init__()\n#         self.block = nn.Sequential(\n#             nn.Linear(dim, dim),\n#             nn.BatchNorm1d(dim),\n#             nn.ReLU(),\n#             nn.Dropout(dropout)\n#         )\n#     def forward(self, x):\n#         return x + self.block(x)\n\n# class AdvancedClassifier(nn.Module):\n#     def __init__(self, input_dim, num_classes):\n#         super().__init__()\n#         # Input Projection\n#         self.input_layer = nn.Sequential(\n#             nn.Linear(input_dim, 512),\n#             nn.BatchNorm1d(512),\n#             nn.ReLU()\n#         )\n        \n#         # Deep Residual Layers (The \"Brain\")\n#         self.res_blocks = nn.Sequential(\n#             ResidualBlock(512),\n#             ResidualBlock(512),\n#             ResidualBlock(512)\n#         )\n        \n#         # Output Projection\n#         self.output_layer = nn.Linear(512, num_classes)\n        \n#     def forward(self, x):\n#         x = self.input_layer(x)\n#         x = self.res_blocks(x)\n#         x = self.output_layer(x)\n#         return x\n\n# # ==========================================\n# # 4. MAIN PIPELINE\n# # ==========================================\n# def main():\n#     print(\"=== STARTING ADVANCED CAFA 6 PIPELINE ===\")\n#     print(f\"DEBUG_MODE: {Config.DEBUG_MODE} (Must be False for submission)\")\n    \n#     # 1. Labels\n#     label_mgr = LabelManager(Config.OBO_FILE, Config.TRAIN_TERMS, Config.NUM_LABELS)\n#     labels_map = label_mgr.get_labels_dict()\n    \n#     # 2. Training Data\n#     limit = 500 if Config.DEBUG_MODE else None\n#     train_ids, train_embeds = extract_embeddings(Config.TRAIN_FASTA, Config.MODEL_NAME, limit=limit)\n    \n#     X_data, y_data = [], []\n#     for pid, emb in zip(train_ids, train_embeds):\n#         if pid in labels_map:\n#             X_data.append(emb)\n#             y_data.append(labels_map[pid])\n            \n#     X_data = np.array(X_data)\n#     y_data = np.array(y_data)\n#     print(f\"Training on {len(X_data)} samples.\")\n    \n#     dataset = ProteinDataset(X_data, y_data)\n#     train_size = int(0.8 * len(dataset))\n#     val_size = len(dataset) - train_size\n#     train_ds, val_ds = random_split(dataset, [train_size, val_size])\n    \n#     train_loader = DataLoader(train_ds, batch_size=Config.BATCH_SIZE, shuffle=True)\n#     val_loader = DataLoader(val_ds, batch_size=Config.BATCH_SIZE)\n    \n#     # 3. Setup Model with Weighted Loss\n#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     model = AdvancedClassifier(Config.EMBED_DIM, Config.NUM_LABELS).to(device)\n    \n#     # IMPROVEMENT: Calculate Positive Weights to fight Class Imbalance\n#     # Simple strategy: roughly scale up positive class error\n#     pos_weight = torch.ones(Config.NUM_LABELS).to(device) * 10.0 \n#     criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n    \n#     optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)\n#     # IMPROVEMENT: Learning Rate Scheduler\n#     scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2, verbose=True)\n    \n#     print(f\"\\n>> Training Advanced Model for {Config.EPOCHS} epochs...\")\n    \n#     for epoch in range(Config.EPOCHS):\n#         model.train()\n#         train_loss = 0\n#         for inputs, targets in train_loader:\n#             inputs, targets = inputs.to(device), targets.to(device)\n#             optimizer.zero_grad()\n#             outputs = model(inputs)\n#             loss = criterion(outputs, targets)\n#             loss.backward()\n#             optimizer.step()\n#             train_loss += loss.item()\n            \n#         model.eval()\n#         val_loss = 0\n#         with torch.no_grad():\n#             for inputs, targets in val_loader:\n#                 inputs, targets = inputs.to(device), targets.to(device)\n#                 outputs = model(inputs)\n#                 loss = criterion(outputs, targets)\n#                 val_loss += loss.item()\n        \n#         avg_val_loss = val_loss / len(val_loader)\n#         print(f\"Epoch {epoch+1} | Train Loss: {train_loss/len(train_loader):.4f} | Val Loss: {avg_val_loss:.4f}\")\n        \n#         # Adjust Learning Rate\n#         scheduler.step(avg_val_loss)\n\n#     torch.save(model.state_dict(), \"cafa6_model_advanced.pth\")\n#     print(\">> Model Saved!\")\n\n#     # 4. Inference\n#     print(\"\\n>> Processing Test Data...\")\n#     test_limit = 200 if Config.DEBUG_MODE else None\n    \n#     # IMPORTANT: The model name MUST match what we used for training\n#     test_ids, test_embeds = extract_embeddings(Config.TEST_FASTA, Config.MODEL_NAME, limit=test_limit)\n    \n#     if len(test_ids) == 0:\n#         print(\"Error: No test data found.\")\n#         return\n\n#     test_dataset = ProteinDataset(test_embeds, labels=None)\n#     test_loader = DataLoader(test_dataset, batch_size=Config.BATCH_SIZE, shuffle=False)\n    \n#     predictions = []\n#     model.eval()\n#     with torch.no_grad():\n#         for inputs in tqdm(test_loader, desc=\"Predicting\"):\n#             inputs = inputs.to(device)\n#             outputs = model(inputs)\n#             probs = torch.sigmoid(outputs).cpu().numpy()\n#             predictions.append(probs)\n            \n#     predictions = np.vstack(predictions)\n    \n#     print(\">> Generating submission.tsv...\")\n#     with open('submission.tsv', 'w') as f:\n#         # NOTE: CAFA usually does NOT want a header line. Writing strictly data.\n#         for i, pid in enumerate(tqdm(test_ids)):\n#             probs = predictions[i]\n#             # Get top 60 predictions\n#             top_indices = np.argsort(probs)[-60:][::-1]\n            \n#             for idx in top_indices:\n#                 score = probs[idx]\n#                 # Lower threshold slightly to ensure we submit enough candidates\n#                 if score > 0.005: \n#                     term_id = label_mgr.idx_to_term[idx]\n#                     f.write(f\"{pid}\\t{term_id}\\t{score:.3f}\\n\")\n                    \n#     print(\"=== PIPELINE COMPLETE ===\")\n\n# if __name__ == \"__main__\":\n#     main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T17:38:20.960992Z","iopub.execute_input":"2025-12-14T17:38:20.961602Z","iopub.status.idle":"2025-12-14T19:52:33.124486Z","shell.execute_reply.started":"2025-12-14T17:38:20.961578Z","shell.execute_reply":"2025-12-14T19:52:33.123536Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms = pd.read_csv(\"/kaggle/working/submission.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"go_term\"])\nprint(train_terms.head())\nprint(\"Total annotations:\", len(train_terms))\nprint(\"Unique GO terms:\", train_terms.go_term.nunique())\nprint(train_terms)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T19:53:28.137544Z","iopub.execute_input":"2025-12-14T19:53:28.137848Z","iopub.status.idle":"2025-12-14T19:53:31.891881Z","shell.execute_reply.started":"2025-12-14T19:53:28.137825Z","shell.execute_reply":"2025-12-14T19:53:31.891038Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import os\n# import torch\n# import torch.nn as nn\n# import torch.optim as optim\n# from torch.utils.data import Dataset, DataLoader, random_split\n# from transformers import AutoTokenizer, EsmModel\n# from Bio import SeqIO\n# import obonet\n# import networkx as nx\n# import pandas as pd\n# import numpy as np\n# from tqdm import tqdm\n# from collections import Counter\n# from sklearn.metrics import f1_score, precision_score, recall_score\n\n# # ==========================================\n# # CONFIGURATION\n# # ==========================================\n# class Config:\n#     # 1. Base Directories\n#     BASE_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n#     WORK_DIR = '/kaggle/working'\n    \n#     # 2. File Paths (Training)\n#     if os.path.exists(os.path.join(BASE_DIR, 'Train')):\n#         TRAIN_FASTA = os.path.join(BASE_DIR, 'Train', 'train_sequences.fasta')\n#         TRAIN_TERMS = os.path.join(BASE_DIR, 'Train', 'train_terms.tsv')\n#         OBO_FILE    = os.path.join(BASE_DIR, 'Train', 'go-basic.obo')\n#     else:\n#         TRAIN_FASTA = os.path.join(BASE_DIR, 'train_sequences.fasta')\n#         TRAIN_TERMS = os.path.join(BASE_DIR, 'train_terms.tsv')\n#         OBO_FILE    = os.path.join(BASE_DIR, 'go-basic.obo')\n\n#     # 3. File Paths (Test)\n#     if os.path.exists(os.path.join(BASE_DIR, 'Test')):\n#         TEST_FASTA  = os.path.join(BASE_DIR, 'Test', 'testsuperset.fasta')\n#     else:\n#         TEST_FASTA  = os.path.join(BASE_DIR, 'testsuperset.fasta')\n\n#     # 4. Model Parameters (Must match training exactly)\n#     MODEL_NAME = \"facebook/esm2_t12_35M_UR50D\" # Updated to match Advanced Training Pipeline\n#     EMBED_DIM = 480     # 480 for 35M model\n#     MAX_SEQ_LEN = 1022  \n#     NUM_LABELS = 2000   \n#     BATCH_SIZE = 32     \n    \n#     # 5. Evaluation Settings\n#     VAL_SPLIT = 0.2     \n#     THRESHOLD = 0.25    # Lower threshold often helps recall in CAFA\n    \n#     # Set to None to process the full dataset (Required for Submission)\n#     # Set to 1000 for quick debugging\n#     LIMIT_SAMPLES = None \n\n# # ==========================================\n# # 1. CLASS DEFINITIONS (Must Match Training)\n# # ==========================================\n# class LabelManager:\n#     def __init__(self, obo_path, terms_path, num_labels=1500):\n#         print(f\">> Loading Ontology...\")\n#         self.graph = obonet.read_obo(obo_path)\n        \n#         # Precompute ancestors\n#         self.term_ancestors = {\n#             term: set(nx.ancestors(self.graph, term)) \n#             for term in tqdm(self.graph.nodes, desc=\"Indexing Ancestors\")\n#         }\n        \n#         print(f\">> Loading Terms...\")\n#         self.terms_df = pd.read_csv(terms_path, sep=\"\\t\")\n        \n#         # Standardize ID column\n#         possible_id_cols = ['EntryID', 'ProteinID', 'UniprotID', 'protein_id']\n#         found_col = None\n#         for col in possible_id_cols:\n#             if col in self.terms_df.columns:\n#                 found_col = col\n#                 break\n#         if found_col:\n#             self.terms_df.rename(columns={found_col: 'protein_id'}, inplace=True)\n#         else:\n#             self.terms_df.rename(columns={self.terms_df.columns[0]: 'protein_id'}, inplace=True)\n            \n#         term_counts = Counter(self.terms_df['term'])\n#         self.selected_terms = [t[0] for t in term_counts.most_common(num_labels)]\n#         self.term_to_idx = {t: i for i, t in enumerate(self.selected_terms)}\n#         self.idx_to_term = {i: t for t, i in self.term_to_idx.items()}\n        \n#         # Pre-compute Parent Indices for Consistency Enforcement\n#         print(\">> Building graph dependency map...\")\n#         self.term_parents_indices = []\n#         for term in self.selected_terms:\n#             parents = []\n#             if term in self.graph:\n#                 # In OBO graph, successors are parents (is_a relationship)\n#                 for p in self.graph.successors(term):\n#                     if p in self.term_to_idx:\n#                         parents.append(self.term_to_idx[p])\n#             self.term_parents_indices.append(parents)\n\n#     def get_labels_dict(self):\n#         protein_groups = self.terms_df.groupby('protein_id')['term'].apply(set)\n#         labels_dict = {}\n#         all_ids = list(protein_groups.index)\n        \n#         if Config.LIMIT_SAMPLES:\n#             all_ids = all_ids[:Config.LIMIT_SAMPLES]\n            \n#         print(\">> Propagating Labels (Full Dataset)...\")\n#         for pid in tqdm(all_ids):\n#             raw_terms = protein_groups[pid]\n#             propagated_terms = set(raw_terms)\n#             for term in raw_terms:\n#                 if term in self.term_ancestors:\n#                     propagated_terms.update(self.term_ancestors[term])\n#             vec = np.zeros(Config.NUM_LABELS, dtype=np.float32)\n#             for term in propagated_terms:\n#                 if term in self.term_to_idx:\n#                     vec[self.term_to_idx[term]] = 1.0\n#             labels_dict[pid] = vec\n#         return labels_dict\n\n#     def enforce_consistency(self, prediction_matrix):\n#         \"\"\"\n#         Propagates scores up the graph.\n#         \"\"\"\n#         consistent_preds = prediction_matrix.copy()\n#         for _ in range(3): # 3 passes is usually sufficient\n#             for i in range(len(self.selected_terms)):\n#                 current_scores = consistent_preds[:, i]\n#                 parents = self.term_parents_indices[i]\n#                 for p_idx in parents:\n#                     consistent_preds[:, p_idx] = np.maximum(\n#                         consistent_preds[:, p_idx], \n#                         current_scores\n#                     )\n#         return consistent_preds\n\n# # --- ADVANCED MODEL ARCHITECTURE ---\n# class ResidualBlock(nn.Module):\n#     def __init__(self, dim, dropout=0.3):\n#         super().__init__()\n#         self.block = nn.Sequential(\n#             nn.Linear(dim, dim),\n#             nn.BatchNorm1d(dim),\n#             nn.ReLU(),\n#             nn.Dropout(dropout)\n#         )\n#     def forward(self, x):\n#         return x + self.block(x)\n\n# class AdvancedClassifier(nn.Module):\n#     def __init__(self, input_dim, num_classes):\n#         super().__init__()\n#         self.input_layer = nn.Sequential(\n#             nn.Linear(input_dim, 512),\n#             nn.BatchNorm1d(512),\n#             nn.ReLU()\n#         )\n#         self.res_blocks = nn.Sequential(\n#             ResidualBlock(512),\n#             ResidualBlock(512),\n#             ResidualBlock(512)\n#         )\n#         self.output_layer = nn.Linear(512, num_classes)\n        \n#     def forward(self, x):\n#         x = self.input_layer(x)\n#         x = self.res_blocks(x)\n#         x = self.output_layer(x)\n#         return x\n\n# def clean_id(header_id):\n#     # Handles \"sp|P12345|NAME\" -> \"P12345\"\n#     # Handles \"A0A0C5B5G6\" -> \"A0A0C5B5G6\"\n#     if \"|\" in header_id:\n#         parts = header_id.split('|')\n#         if len(parts) >= 2: return parts[1]\n#     return header_id\n\n# def extract_embeddings(fasta_path, model_name, limit=None):\n#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     print(f\">> Extracting Features from {os.path.basename(fasta_path)}...\")\n#     tokenizer = AutoTokenizer.from_pretrained(model_name)\n#     model = EsmModel.from_pretrained(model_name).to(device)\n#     model.eval()\n    \n#     seqs = []\n#     ids = []\n#     # Using 'fasta' format from BioPython\n#     for i, record in enumerate(SeqIO.parse(fasta_path, \"fasta\")):\n#         if limit and i >= limit: break\n#         ids.append(clean_id(str(record.id)))\n#         seqs.append(str(record.seq)[:Config.MAX_SEQ_LEN]) \n        \n#     embeddings = []\n#     with torch.no_grad():\n#         for i in tqdm(range(0, len(seqs), Config.BATCH_SIZE)):\n#             batch_seqs = seqs[i : i + Config.BATCH_SIZE]\n#             inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True)\n#             inputs = {k: v.to(device) for k, v in inputs.items()}\n#             outputs = model(**inputs)\n#             embeddings.append(outputs.last_hidden_state[:, 0, :].cpu().numpy())\n            \n#     if not embeddings: return [], np.array([])\n#     return ids, np.vstack(embeddings)\n\n# class ProteinDataset(Dataset):\n#     def __init__(self, embeddings, labels=None):\n#         self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n#         self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n#     def __len__(self): return len(self.embeddings)\n#     def __getitem__(self, idx): \n#         if self.labels is not None: return self.embeddings[idx], self.labels[idx]\n#         return self.embeddings[idx]\n\n# # ==========================================\n# # 2. EVALUATION & TEST LOGIC\n# # ==========================================\n# def evaluate_and_test():\n#     print(\"=== STARTING EVALUATION & TESTING ===\")\n    \n#     # 1. Setup Label Manager\n#     label_mgr = LabelManager(Config.OBO_FILE, Config.TRAIN_TERMS, Config.NUM_LABELS)\n    \n#     # 2. Load Model\n#     device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n#     model = AdvancedClassifier(Config.EMBED_DIM, Config.NUM_LABELS).to(device)\n    \n#     # Check possible paths for the trained model\n#     possible_paths = [\n#         os.path.join(Config.WORK_DIR, 'cafa6_model_advanced.pth'),\n#         'cafa6_model_advanced.pth'\n#     ]\n    \n#     model_path = None\n#     for path in possible_paths:\n#         if os.path.exists(path):\n#             model_path = path\n#             break\n            \n#     if model_path:\n#         model.load_state_dict(torch.load(model_path, map_location=device))\n#         print(f\">> Loaded trained model from: {model_path}\")\n#     else:\n#         print(\"!! ERROR: Model file not found. Please run training pipeline first!\")\n#         return\n\n#     # --- PART A: VALIDATION (Evaluate Metrics) ---\n#     print(\"\\n[PART A] Evaluating Model Performance on Training Split...\")\n#     # Only need labels for validation\n#     labels_map = label_mgr.get_labels_dict()\n    \n#     train_ids, train_embeds = extract_embeddings(Config.TRAIN_FASTA, Config.MODEL_NAME, limit=Config.LIMIT_SAMPLES)\n    \n#     X_val, y_val = [], []\n#     valid_count = 0\n#     for pid, emb in zip(train_ids, train_embeds):\n#         if pid in labels_map:\n#             X_val.append(emb)\n#             y_val.append(labels_map[pid])\n#             valid_count += 1\n            \n#     print(f\"Matched {valid_count} proteins for validation.\")\n    \n#     if valid_count > 0:\n#         X_val = np.array(X_val)\n#         y_val = np.array(y_val)\n        \n#         # Use a subset for validation (simulate held-out data)\n#         dataset = ProteinDataset(X_val, y_val)\n#         val_len = int(len(dataset) * Config.VAL_SPLIT)\n#         if val_len == 0: val_len = len(dataset) # Fallback for small tests\n#         _, val_ds = random_split(dataset, [len(dataset) - val_len, val_len])\n        \n#         val_loader = DataLoader(val_ds, batch_size=Config.BATCH_SIZE, shuffle=False)\n        \n#         model.eval()\n#         all_preds, all_targets = [], []\n        \n#         with torch.no_grad():\n#             for inputs, targets in val_loader:\n#                 inputs = inputs.to(device)\n#                 outputs = model(inputs)\n#                 probs = torch.sigmoid(outputs).cpu().numpy()\n#                 all_preds.append(probs)\n#                 all_targets.append(targets.numpy())\n                \n#         all_preds = np.vstack(all_preds)\n#         all_targets = np.vstack(all_targets)\n        \n#         # Enforce consistency before scoring\n#         all_preds = label_mgr.enforce_consistency(all_preds)\n#         binary_preds = (all_preds > Config.THRESHOLD).astype(int)\n        \n#         micro_f1 = f1_score(all_targets, binary_preds, average='micro', zero_division=0)\n#         print(f\"\\n>> Validation Micro F1 Score: {micro_f1:.4f}\")\n#     else:\n#         print(\"Skipping validation (no matched labels found or Debug Limit too low).\")\n\n#     # --- PART B: TEST INFERENCE (Generate Submission) ---\n#     print(\"\\n[PART B] Running Inference on Test Dataset...\")\n#     print(f\"Reading Test File: {Config.TEST_FASTA}\")\n    \n#     test_ids, test_embeds = extract_embeddings(Config.TEST_FASTA, Config.MODEL_NAME, limit=Config.LIMIT_SAMPLES)\n    \n#     if len(test_ids) == 0:\n#         print(\"No test sequences found.\")\n#         return\n\n#     test_dataset = ProteinDataset(test_embeds, labels=None)\n#     test_loader = DataLoader(test_dataset, batch_size=Config.BATCH_SIZE, shuffle=False)\n    \n#     test_probs = []\n#     with torch.no_grad():\n#         for inputs in tqdm(test_loader, desc=\"Predicting Test\"):\n#             inputs = inputs.to(device)\n#             outputs = model(inputs)\n#             probs = torch.sigmoid(outputs).cpu().numpy()\n#             test_probs.append(probs)\n            \n#     test_probs = np.vstack(test_probs)\n    \n#     print(\">> Enforcing Consistency on Test Predictions...\")\n#     test_probs = label_mgr.enforce_consistency(test_probs)\n    \n#     # Generate Output File\n#     out_file = 'submission.tsv'\n#     print(f\">> Writing predictions to {out_file}...\")\n    \n#     with open(out_file, 'w') as f:\n#         # Standard CAFA format: ProteinID <tab> GO_Term <tab> Score\n#         # No header\n#         for i, pid in enumerate(tqdm(test_ids)):\n#             probs = test_probs[i]\n#             # Get top 60 predictions to ensure we cover likely terms\n#             top_indices = np.argsort(probs)[-60:][::-1]\n            \n#             for idx in top_indices:\n#                 score = probs[idx]\n#                 if score > 0.001: # Very low threshold to ensure coverage\n#                     term_id = label_mgr.idx_to_term[idx]\n#                     f.write(f\"{pid}\\t{term_id}\\t{score:.3f}\\n\")\n                \n#     print(f\"\\n>> SUCCESS! '{out_file}' is ready for submission.\")\n\n# if __name__ == \"__main__\":\n#     evaluate_and_test()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T20:05:29.814927Z","iopub.execute_input":"2025-12-14T20:05:29.815766Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# CAFA 6 Hybrid Ensemble Pipeline","metadata":{}},{"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom transformers import AutoTokenizer, EsmModel\nfrom Bio import SeqIO\nimport obonet\nimport networkx as nx\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nfrom collections import Counter\nimport gc\n\n# ==========================================\n# CONFIGURATION\n# ==========================================\nclass Config:\n    # 1. Paths\n    BASE_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n    if os.path.exists(os.path.join(BASE_DIR, 'Train')):\n        TRAIN_FASTA = os.path.join(BASE_DIR, 'Train', 'train_sequences.fasta')\n        TRAIN_TERMS = os.path.join(BASE_DIR, 'Train', 'train_terms.tsv')\n        OBO_FILE    = os.path.join(BASE_DIR, 'Train', 'go-basic.obo')\n    else:\n        TRAIN_FASTA = os.path.join(BASE_DIR, 'train_sequences.fasta')\n        TRAIN_TERMS = os.path.join(BASE_DIR, 'train_terms.tsv')\n        OBO_FILE    = os.path.join(BASE_DIR, 'go-basic.obo')\n\n    if os.path.exists(os.path.join(BASE_DIR, 'Test')):\n        TEST_FASTA  = os.path.join(BASE_DIR, 'Test', 'testsuperset.fasta')\n    else:\n        TEST_FASTA  = os.path.join(BASE_DIR, 'testsuperset.fasta')\n    \n    # 2. Model Parameters\n    MODEL_NAME = \"facebook/esm2_t12_35M_UR50D\" \n    EMBED_DIM = 480     \n    \n    MAX_SEQ_LEN = 1022  \n    NUM_LABELS = 2000   \n    BATCH_SIZE = 64     # Increased to 64 for speed (Mixed Precision allows this)\n    LEARNING_RATE = 5e-4\n    \n    # Increased to 10 (Safe with AMP optimization)\n    EPOCHS = 100        \n    \n    # 3. Ensemble Weights \n    WEIGHT_NN = 0.5\n    WEIGHT_KNN = 0.5\n    \n    DEBUG_MODE = False \n\n# ==========================================\n# 1. ONTOLOGY & LABELS\n# ==========================================\nclass LabelManager:\n    def __init__(self, obo_path, terms_path, num_labels=1500):\n        print(f\">> Loading Ontology...\")\n        self.graph = obonet.read_obo(obo_path)\n        self.term_ancestors = {\n            term: set(nx.ancestors(self.graph, term)) \n            for term in tqdm(self.graph.nodes, desc=\"Indexing Ancestors\")\n        }\n        \n        print(f\">> Loading Terms...\")\n        self.terms_df = pd.read_csv(terms_path, sep=\"\\t\")\n        \n        # ID Standardization\n        possible_id_cols = ['EntryID', 'ProteinID', 'UniprotID', 'protein_id']\n        found_col = None\n        for col in possible_id_cols:\n            if col in self.terms_df.columns:\n                found_col = col\n                break\n        if found_col:\n            self.terms_df.rename(columns={found_col: 'protein_id'}, inplace=True)\n        else:\n            self.terms_df.rename(columns={self.terms_df.columns[0]: 'protein_id'}, inplace=True)\n            \n        term_counts = Counter(self.terms_df['term'])\n        self.selected_terms = [t[0] for t in term_counts.most_common(num_labels)]\n        self.term_to_idx = {t: i for i, t in enumerate(self.selected_terms)}\n        self.idx_to_term = {i: t for t, i in self.term_to_idx.items()}\n\n    def get_labels_dict(self):\n        protein_groups = self.terms_df.groupby('protein_id')['term'].apply(set)\n        labels_dict = {}\n        all_ids = list(protein_groups.index)\n        \n        if Config.DEBUG_MODE: all_ids = all_ids[:500]\n            \n        print(\">> Propagating Labels...\")\n        for pid in tqdm(all_ids):\n            raw_terms = protein_groups[pid]\n            propagated_terms = set(raw_terms)\n            for term in raw_terms:\n                if term in self.term_ancestors:\n                    propagated_terms.update(self.term_ancestors[term])\n            \n            vec = np.zeros(Config.NUM_LABELS, dtype=np.float32)\n            for term in propagated_terms:\n                if term in self.term_to_idx:\n                    vec[self.term_to_idx[term]] = 1.0\n            labels_dict[pid] = vec\n        return labels_dict\n\n# ==========================================\n# 2. OPTIMIZED FEATURE EXTRACTION\n# ==========================================\ndef clean_id(header_id):\n    if \"|\" in header_id:\n        parts = header_id.split('|')\n        if len(parts) >= 2: return parts[1]\n    return header_id\n\ndef extract_embeddings(fasta_path, model_name, limit=None):\n    \"\"\"\n    Optimized with Smart Batching (Sort by length) and AMP (Mixed Precision)\n    \"\"\"\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\">> Extracting Features from {os.path.basename(fasta_path)}...\")\n    \n    tokenizer = AutoTokenizer.from_pretrained(model_name)\n    model = EsmModel.from_pretrained(model_name).to(device)\n    model.eval()\n    \n    # 1. Read all sequences first to sort them\n    print(\"   Reading FASTA...\")\n    indexed_seqs = [] # (original_index, clean_id, sequence)\n    \n    for i, record in enumerate(SeqIO.parse(fasta_path, \"fasta\")):\n        if limit and i >= limit: break\n        s_str = str(record.seq)[:Config.MAX_SEQ_LEN]\n        indexed_seqs.append((i, clean_id(str(record.id)), s_str))\n    \n    if not indexed_seqs:\n        return [], np.array([])\n\n    # 2. Smart Batching: Sort by sequence length\n    # This drastically reduces padding and speeds up inference\n    print(\"   Sorting sequences for Smart Batching...\")\n    indexed_seqs.sort(key=lambda x: len(x[2]))\n    \n    sorted_indices = [x[0] for x in indexed_seqs]\n    sorted_ids = [x[1] for x in indexed_seqs]\n    sorted_seqs = [x[2] for x in indexed_seqs]\n    \n    embeddings = []\n    # Increase inference batch size slightly\n    batch_size = 64\n    \n    # 3. Inference with AMP (Automatic Mixed Precision)\n    print(\"   Running Inference...\")\n    with torch.no_grad():\n        for i in tqdm(range(0, len(sorted_seqs), batch_size), desc=\"Embed Batches\"):\n            batch_seqs = sorted_seqs[i : i + batch_size]\n            \n            inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True)\n            inputs = {k: v.to(device) for k, v in inputs.items()}\n            \n            # Autocast for speed\n            with torch.cuda.amp.autocast():\n                outputs = model(**inputs)\n                # Take CLS token\n                batch_emb = outputs.last_hidden_state[:, 0, :]\n            \n            embeddings.append(batch_emb.cpu().numpy().astype(np.float16)) # Save RAM\n            \n    all_embeddings = np.vstack(embeddings)\n    \n    # 4. Restore Original Order\n    # For KNN mapping, we usually want sorted to match embeddings, but IDs must track.\n    # Returning sorted_ids matches the embeddings order.\n    return sorted_ids, all_embeddings\n\n# ==========================================\n# 3. AGENT A: DEEP LEARNING MODEL\n# ==========================================\nclass ProteinDataset(Dataset):\n    def __init__(self, embeddings, labels=None):\n        self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n        self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n    def __len__(self): return len(self.embeddings)\n    def __getitem__(self, idx):\n        if self.labels is not None: return self.embeddings[idx], self.labels[idx]\n        return self.embeddings[idx]\n\nclass ResNetClassifier(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super().__init__()\n        self.input = nn.Sequential(nn.Linear(input_dim, 512), nn.BatchNorm1d(512), nn.ReLU())\n        self.res1 = nn.Sequential(nn.Linear(512, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3))\n        self.res2 = nn.Sequential(nn.Linear(512, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3))\n        self.out = nn.Linear(512, num_classes)\n        \n    def forward(self, x):\n        x = self.input(x)\n        x = x + self.res1(x) # Residual connection\n        x = x + self.res2(x) # Residual connection\n        return self.out(x)\n\n# ==========================================\n# 4. AGENT B: KNN / HOMOLOGY (Optimized)\n# ==========================================\ndef run_knn_inference(train_embeds, train_labels, test_embeds, k=5):\n    print(f\">> Running KNN (Homology Transfer) with k={k}...\")\n    \n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    \n    # Use float16 for massive speedup and memory savings\n    X_train = torch.tensor(train_embeds, dtype=torch.float16).to(device)\n    Y_train = torch.tensor(train_labels, dtype=torch.float32).to(device) # Labels need float32 for averaging\n    X_test = torch.tensor(test_embeds, dtype=torch.float16).to(device)\n    \n    # Normalize\n    X_train = torch.nn.functional.normalize(X_train, p=2, dim=1)\n    X_test = torch.nn.functional.normalize(X_test, p=2, dim=1)\n    \n    num_test = X_test.shape[0]\n    num_labels = Y_train.shape[1]\n    knn_preds = np.zeros((num_test, num_labels), dtype=np.float32)\n    \n    chunk_size = 1000\n    \n    with torch.no_grad():\n        for i in tqdm(range(0, num_test, chunk_size), desc=\"KNN Search\"):\n            test_chunk = X_test[i : i + chunk_size]\n            \n            # Matrix Multiplication in FP16\n            sim_matrix = torch.matmul(test_chunk, X_train.T)\n            \n            top_vals, top_inds = torch.topk(sim_matrix, k=k, dim=1)\n            \n            # Convert back to float32 for calculations\n            weights = torch.relu(top_vals).float()\n            weights = weights / (weights.sum(dim=1, keepdim=True) + 1e-6)\n            \n            neighbor_labels = Y_train[top_inds]\n            weighted_preds = (weights.unsqueeze(2) * neighbor_labels).sum(dim=1)\n            \n            knn_preds[i : i + chunk_size] = weighted_preds.cpu().numpy()\n            \n            del sim_matrix, top_vals, top_inds, neighbor_labels, weighted_preds\n            torch.cuda.empty_cache()\n            \n    return knn_preds\n\n# ==========================================\n# 5. MAIN EXECUTION PIPELINE\n# ==========================================\ndef main():\n    print(\"=== STARTING OPTIMIZED CAFA 6 HYBRID PIPELINE ===\")\n    \n    # 1. Setup Labels\n    label_mgr = LabelManager(Config.OBO_FILE, Config.TRAIN_TERMS, Config.NUM_LABELS)\n    labels_map = label_mgr.get_labels_dict()\n    \n    # 2. Extract Training Features\n    limit = 500 if Config.DEBUG_MODE else None\n    train_ids, train_embeds = extract_embeddings(Config.TRAIN_FASTA, Config.MODEL_NAME, limit=limit)\n    \n    # Align Data\n    X_train_raw, y_train_raw = [], []\n    print(\">> Aligning Training Data...\")\n    for i, pid in enumerate(train_ids):\n        if pid in labels_map:\n            X_train_raw.append(train_embeds[i])\n            y_train_raw.append(labels_map[pid])\n            \n    X_train = np.array(X_train_raw)\n    y_train = np.array(y_train_raw)\n    print(f\"Training on {len(X_train)} samples.\")\n    \n    # 3. Train Neural Network (Agent A)\n    dataset = ProteinDataset(X_train, y_train)\n    train_size = int(0.9 * len(dataset))\n    train_ds, val_ds = random_split(dataset, [train_size, len(dataset) - train_size])\n    train_loader = DataLoader(train_ds, batch_size=Config.BATCH_SIZE, shuffle=True)\n    \n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = ResNetClassifier(Config.EMBED_DIM, Config.NUM_LABELS).to(device)\n    \n    pos_weight = torch.ones(Config.NUM_LABELS).to(device) * 12.0\n    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n    optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)\n    \n    # Added Scaler for Mixed Precision Training (Faster!)\n    scaler = torch.cuda.amp.GradScaler()\n    \n    print(f\"\\n>> Training Agent A (Neural Net) for {Config.EPOCHS} epochs...\")\n    for epoch in range(Config.EPOCHS):\n        model.train()\n        total_loss = 0\n        for inputs, targets in train_loader:\n            inputs, targets = inputs.to(device), targets.to(device)\n            optimizer.zero_grad()\n            \n            # Autocast for faster training\n            with torch.cuda.amp.autocast():\n                outputs = model(inputs)\n                loss = criterion(outputs, targets)\n            \n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            \n            total_loss += loss.item()\n        print(f\"Epoch {epoch+1} Loss: {total_loss/len(train_loader):.4f}\")\n        \n    # 4. Process Test Data\n    print(\"\\n>> Processing Test Data (Optimized)...\")\n    test_limit = 200 if Config.DEBUG_MODE else None\n    test_ids, test_embeds = extract_embeddings(Config.TEST_FASTA, Config.MODEL_NAME, limit=test_limit)\n    \n    if len(test_ids) == 0: return\n\n    # 5. AGENT A Prediction\n    print(\"\\n>> Agent A Inference...\")\n    test_ds = ProteinDataset(test_embeds)\n    test_loader = DataLoader(test_ds, batch_size=Config.BATCH_SIZE, shuffle=False)\n    nn_preds = []\n    model.eval()\n    with torch.no_grad():\n        for inputs in tqdm(test_loader):\n            inputs = inputs.to(device)\n            # Mixed precision inference\n            with torch.cuda.amp.autocast():\n                out = model(inputs)\n                nn_preds.append(torch.sigmoid(out).float().cpu().numpy())\n    nn_preds = np.vstack(nn_preds)\n    \n    # 6. AGENT B Prediction\n    print(\"\\n>> Agent B Inference (KNN)...\")\n    knn_preds = run_knn_inference(X_train, y_train, test_embeds, k=5)\n    \n    # 7. ENSEMBLE\n    print(\"\\n>> Ensembling...\")\n    final_preds = (Config.WEIGHT_NN * nn_preds) + (Config.WEIGHT_KNN * knn_preds)\n    \n    # 8. Generate Submission\n    print(\">> Generating submission.tsv...\")\n    with open('submission.tsv', 'w') as f:\n        for i, pid in enumerate(tqdm(test_ids)):\n            probs = final_preds[i]\n            top_indices = np.argsort(probs)[-60:][::-1]\n            for idx in top_indices:\n                score = probs[idx]\n                if score > 0.001:\n                    term_id = label_mgr.idx_to_term[idx]\n                    f.write(f\"{pid}\\t{term_id}\\t{score:.3f}\\n\")\n                    \n    print(\"=== PIPELINE COMPLETE ===\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-17T04:06:07.652551Z","iopub.execute_input":"2025-12-17T04:06:07.652867Z","iopub.status.idle":"2025-12-17T05:26:28.779178Z","shell.execute_reply.started":"2025-12-17T04:06:07.652846Z","shell.execute_reply":"2025-12-17T05:26:28.778491Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}