{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":6247561,"sourceType":"datasetVersion","datasetId":3590060},{"sourceId":14292632,"sourceType":"datasetVersion","datasetId":9123399}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"\n

\nšŸ§¬ Predicting Protein Function \n

\n\n

\n šŸ“Œ What is the task?\n

\n

\n This competition asks us to predict Gene Ontology (GO) terms for proteins based on their sequences. GO terms are standardized labels that describe:\n\n

\n

\n A single protein can have multiple functions, so this is a **multi-label classification** problem. We're not picking one category ā€” we're predicting all the relevant functional labels simultaneously.\n

\n

\n\n\n \n\n\n\n\n","metadata":{}},{"cell_type":"code","source":"!pip install obonet ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:16:37.911976Z","iopub.execute_input":"2026-01-13T12:16:37.912739Z","iopub.status.idle":"2026-01-13T12:16:41.188220Z","shell.execute_reply.started":"2026-01-13T12:16:37.912711Z","shell.execute_reply":"2026-01-13T12:16:41.187453Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 1. Importing Libraries\n

\n","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport obonet\nimport torch\nfrom torch import nn\nimport networkx as nx\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchmetrics.classification import MultilabelF1Score\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom matplotlib import pyplot as plt\nfrom tqdm import tqdm\nimport warnings\n\nwarnings.filterwarnings(\"ignore\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:16:41.190211Z","iopub.execute_input":"2026-01-13T12:16:41.190532Z","iopub.status.idle":"2026-01-13T12:16:41.195846Z","shell.execute_reply.started":"2026-01-13T12:16:41.190505Z","shell.execute_reply":"2026-01-13T12:16:41.195185Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 2. Configuration\n

\n","metadata":{}},{"cell_type":"code","source":"class Config:\n \"\"\"Central configuration for all hyperparameters and paths\"\"\"\n \n # Data paths\n train_embeds_path = \"/kaggle/input/cafa-6-protbert-embeddings-dataset/train_embeds_final.npy\"\n test_embeds_path = \"/kaggle/input/cafa-6-protbert-embeddings-dataset/test_embeds_final.npy\"\n train_ids_path = \"/kaggle/input/cafa-6-protbert-embeddings-dataset/train_ids_final.npy\"\n test_ids_path = \"/kaggle/input/cafa-6-protbert-embeddings-dataset/test_ids_final.npy\"\n train_terms_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n obo_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n\n # Model architecture\n input_dim = 1024 \n num_labels = 200 \n hidden_dims = [1024, 512, 512] \n dropout = 0.3 \n\n # Training parameters\n batch_size = 128\n train_split = 0.9 \n lr = 1e-3 \n epochs = 100\n patience = 15 \n\n # Device configuration\n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n seed = 42\n\ncfg = Config()\nprint(f\"Using device: {cfg.device}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:16:41.196940Z","iopub.execute_input":"2026-01-13T12:16:41.197433Z","iopub.status.idle":"2026-01-13T12:16:41.212011Z","shell.execute_reply.started":"2026-01-13T12:16:41.197411Z","shell.execute_reply":"2026-01-13T12:16:41.211296Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 3. Data Loading\n

\n","metadata":{}},{"cell_type":"code","source":"print(\"\\n\" + \"=\"*60)\nprint(\"LOADING DATA\")\nprint(\"=\"*60)\n\n# Load pre-computed ProtBERT embeddings\ntrain_embeds = np.load(cfg.train_embeds_path)\ntest_embeds = np.load(cfg.test_embeds_path)\ntrain_ids = np.load(cfg.train_ids_path)\ntest_ids = np.load(cfg.test_ids_path)\n\nprint(f\"Train embeddings shape: {train_embeds.shape}\")\nprint(f\"Test embeddings shape: {test_embeds.shape}\")\n\n# Load GO term labels\ntrain_terms_df = pd.read_csv(cfg.train_terms_path, sep=\"\\t\")\n\n# Group GO terms by protein ID\nlabels_df = (\n train_terms_df\n .groupby(\"EntryID\")[\"term\"]\n .apply(list)\n .reset_index(name=\"labels\")\n)\n\n# Select top N most frequent GO terms\ntop_go = (\n train_terms_df[\"term\"]\n .value_counts()\n .iloc[:cfg.num_labels]\n .index\n .tolist()\n)\n\nprint(f\"\\nUnique proteins with labels: {len(labels_df)}\")\nprint(f\"Top GO terms selected: {len(top_go)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:16:41.213820Z","iopub.execute_input":"2026-01-13T12:16:41.214070Z","iopub.status.idle":"2026-01-13T12:16:43.160952Z","shell.execute_reply.started":"2026-01-13T12:16:41.214029Z","shell.execute_reply":"2026-01-13T12:16:43.160213Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 4. Data Preprocessing\n

\n","metadata":{}},{"cell_type":"code","source":"go_graph = obonet.read_obo(cfg.obo_path)\n\ndef get_ancestors(graph, term):\n if term not in graph:\n return set()\n return nx.descendants(graph, term) | {term}\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"PROPAGATING GO TERM ANCESTORS\")\nprint(\"=\"*60)\n\n# Select top N most frequent GO terms\ntop_go = (\n train_terms_df[\"term\"]\n .value_counts()\n .iloc[:cfg.num_labels]\n .index\n .tolist()\n)\ntop_go_set = set(top_go)\n\n# Group GO terms by protein ID\nprotein_to_terms = train_terms_df.groupby(\"EntryID\")[\"term\"].apply(set).to_dict()\n\n# Propagate ancestors for each protein\npropagated_labels = {}\nprint(\"Propagating ancestors for each protein...\")\n\nfor protein_id, terms in tqdm(protein_to_terms.items(), desc=\"Propagating\"):\n all_ancestors = set()\n for term in terms:\n # Get all ancestors for this term\n ancestors = get_ancestors(go_graph, term)\n all_ancestors.update(ancestors)\n \n # Filter to only keep terms in our top N list\n propagated_labels[protein_id] = all_ancestors.intersection(top_go_set)\n\n# Create labels DataFrame with propagated terms\nlabels_data = []\nfor protein_id, terms in propagated_labels.items():\n if terms: # Only include proteins with at least one term\n labels_data.append({\n \"EntryID\": protein_id,\n \"labels\": list(terms)\n })\n\nlabels_df = pd.DataFrame(labels_data)\n\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"PREPROCESSING DATA\")\nprint(\"=\"*60)\n\n# Initialize MultiLabelBinarizer for one-hot encoding GO terms\nmlb = MultiLabelBinarizer(classes=top_go)\nmlb.fit([top_go])\n\n# Convert GO term lists to binary target vectors\nlabels_df[\"targets\"] = list(mlb.transform(labels_df[\"labels\"]))\n\n# Create training DataFrame with embeddings and labels\ntrain_emb_df = pd.DataFrame({\n \"EntryID\": train_ids,\n \"embed\": list(train_embeds)\n})\n\ntrain_df = train_emb_df.merge(labels_df, on=\"EntryID\", how=\"inner\")\n\n# Create test DataFrame\ntest_df = pd.DataFrame({\n \"EntryID\": test_ids,\n \"embed\": list(test_embeds)\n})\n\nprint(f\"Training data shape: {train_df.shape}\")\nprint(f\"Test data shape: {test_df.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:16:43.162028Z","iopub.execute_input":"2026-01-13T12:16:43.162426Z","iopub.status.idle":"2026-01-13T12:17:01.359540Z","shell.execute_reply.started":"2026-01-13T12:16:43.162405Z","shell.execute_reply":"2026-01-13T12:17:01.358776Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 5. Model Definition\n

\n","metadata":{}},{"cell_type":"code","source":"class MLP(nn.Module):\n \n def __init__(self, input_dim=1024, num_labels=500, hidden_dims=[1024, 512, 512], dropout=0.3):\n super().__init__()\n self.layers = nn.ModuleList()\n self.bns = nn.ModuleList()\n self.dropouts = nn.ModuleList()\n\n # Build hidden layers\n prev_dim = input_dim\n for h in hidden_dims:\n self.layers.append(nn.Linear(prev_dim, h))\n self.bns.append(nn.BatchNorm1d(h))\n self.dropouts.append(nn.Dropout(dropout))\n prev_dim = h\n\n # Output layer\n self.out = nn.Linear(prev_dim, num_labels)\n\n def forward(self, x):\n residual = x\n \n # Pass through hidden layers with residual connections\n for i, layer in enumerate(self.layers):\n x = layer(x)\n x = self.bns[i](x)\n x = torch.relu(x)\n x = self.dropouts[i](x)\n \n # Add residual connection every 2 layers\n if i % 2 == 1:\n if residual.shape[1] == x.shape[1]:\n x = x + residual\n residual = x\n else:\n residual = x\n \n # Output layer\n x = self.out(x)\n return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:17:01.360504Z","iopub.execute_input":"2026-01-13T12:17:01.360766Z","iopub.status.idle":"2026-01-13T12:17:01.367783Z","shell.execute_reply.started":"2026-01-13T12:17:01.360745Z","shell.execute_reply":"2026-01-13T12:17:01.367202Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 6. Dataset and DataLoaders\n

\n","metadata":{}},{"cell_type":"code","source":"class ProteinFunctionDataset(Dataset):\n \"\"\"Custom Dataset for protein embeddings and GO term labels\"\"\"\n \n def __init__(self, df, train=True):\n self.df = df.reset_index(drop=True)\n self.train = train\n\n def __len__(self):\n return len(self.df)\n\n def __getitem__(self, idx):\n embed = torch.tensor(self.df.loc[idx, \"embed\"], dtype=torch.float32)\n \n if self.train:\n target = torch.tensor(self.df.loc[idx, \"targets\"], dtype=torch.float32)\n return embed, target\n else:\n entry_id = self.df.loc[idx, \"EntryID\"]\n return embed, entry_id\n\n# Create datasets\nfull_train_dataset = ProteinFunctionDataset(train_df, train=True)\ntest_dataset = ProteinFunctionDataset(test_df, train=False)\n\n# Split training data into train and validation sets\ntrain_size = int(len(full_train_dataset) * cfg.train_split)\nval_size = len(full_train_dataset) - train_size\n\ng = torch.Generator().manual_seed(cfg.seed)\ntrain_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size], generator=g)\n\n# Create dataloaders\ntrain_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False)\ntest_loader = DataLoader(test_dataset, batch_size=cfg.batch_size, shuffle=False)\n\nprint(f\"\\nTrain batches: {len(train_loader)}\")\nprint(f\"Validation batches: {len(val_loader)}\")\nprint(f\"Test batches: {len(test_loader)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:17:01.368662Z","iopub.execute_input":"2026-01-13T12:17:01.368944Z","iopub.status.idle":"2026-01-13T12:17:01.499280Z","shell.execute_reply.started":"2026-01-13T12:17:01.368923Z","shell.execute_reply":"2026-01-13T12:17:01.498657Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 7. Training loop\n

\n","metadata":{}},{"cell_type":"code","source":"def train_model(model, train_loader, val_loader, cfg, save_path=\"best_model.pt\"):\n \n \n model.to(cfg.device)\n\n # Initialize optimizer, scheduler, and loss function\n optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)\n scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n optimizer, mode=\"min\", factor=0.5, patience=2\n )\n loss_fn = nn.CrossEntropyLoss()\n # loss_fn = nn.BCEWithLogitsLoss() \n f1_score = MultilabelF1Score(num_labels=cfg.num_labels).to(cfg.device)\n\n # Tracking metrics\n history = {\n \"train_loss\": [],\n \"val_loss\": [],\n \"train_f1\": [],\n \"val_f1\": []\n }\n\n best_val_loss = float(\"inf\")\n patience_counter = 0\n\n print(\"\\n\" + \"=\"*60)\n print(\"TRAINING\")\n print(\"=\"*60)\n\n for epoch in range(cfg.epochs):\n # ===== TRAINING PHASE =====\n model.train()\n train_losses, train_f1s = [], []\n\n for embeds, targets in train_loader:\n embeds = embeds.to(cfg.device)\n targets = targets.to(cfg.device)\n\n # Forward pass\n optimizer.zero_grad()\n logits = model(embeds)\n loss = loss_fn(logits, targets)\n \n # Backward pass\n loss.backward()\n optimizer.step()\n\n # Calculate metrics\n f1 = f1_score(logits, targets)\n train_losses.append(loss.item())\n train_f1s.append(f1.item())\n\n avg_train_loss = np.mean(train_losses)\n avg_train_f1 = np.mean(train_f1s)\n\n # ===== VALIDATION PHASE =====\n model.eval()\n val_losses, val_f1s = [], []\n\n with torch.no_grad():\n for embeds, targets in val_loader:\n embeds = embeds.to(cfg.device)\n targets = targets.to(cfg.device)\n\n logits = model(embeds)\n loss = loss_fn(logits, targets)\n\n f1 = f1_score(logits, targets)\n val_losses.append(loss.item())\n val_f1s.append(f1.item())\n\n avg_val_loss = np.mean(val_losses)\n avg_val_f1 = np.mean(val_f1s)\n\n # Save metrics\n history[\"train_loss\"].append(avg_train_loss)\n history[\"val_loss\"].append(avg_val_loss)\n history[\"train_f1\"].append(avg_train_f1)\n history[\"val_f1\"].append(avg_val_f1)\n\n # Print progress\n print(\n f\"Epoch {epoch+1:3d}/{cfg.epochs} | \"\n f\"Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | \"\n f\"Train F1: {avg_train_f1:.4f} | Val F1: {avg_val_f1:.4f}\"\n )\n\n # Learning rate scheduling\n scheduler.step(avg_val_loss)\n\n # ===== CHECKPOINTING & EARLY STOPPING =====\n if avg_val_loss < best_val_loss:\n best_val_loss = avg_val_loss\n patience_counter = 0\n torch.save(model.state_dict(), save_path)\n print(f\"šŸ’¾ Best model saved at epoch {epoch+1}\")\n else:\n patience_counter += 1\n print(f\"ā³ No improvement. Patience: {patience_counter}/{cfg.patience}\")\n\n if patience_counter >= cfg.patience:\n print(\"šŸ›‘ Early stopping triggered.\")\n break\n\n print(f\"\\nāœ… Training complete! Best Val Loss: {best_val_loss:.4f}\")\n return history\n\n\nmodel = MLP(\n input_dim=cfg.input_dim,\n num_labels=cfg.num_labels,\n hidden_dims=cfg.hidden_dims,\n dropout=cfg.dropout\n)\n\nhistory = train_model(model, train_loader, val_loader, cfg, save_path=\"best_model.pt\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:17:01.500251Z","iopub.execute_input":"2026-01-13T12:17:01.500481Z","iopub.status.idle":"2026-01-13T12:20:03.100221Z","shell.execute_reply.started":"2026-01-13T12:17:01.500461Z","shell.execute_reply":"2026-01-13T12:20:03.099407Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 8. Probert Submission \n

\n","metadata":{}},{"cell_type":"code","source":"# Load the best model checkpoint saved during training\nmodel.load_state_dict(torch.load(\"best_model.pt\", map_location=cfg.device))\nprint(\"āœ… Best model loaded successfully!\")\n\ndef generate_submission(\n model,\n test_loader,\n go_terms,\n device,\n output_path=\"submission.tsv\",\n max_terms_per_protein=1500,\n min_prob=1e-6\n):\n \n \n model.eval()\n \n all_ids = []\n all_probs = []\n\n # Generate predictions\n with torch.no_grad():\n for embeds, ids in tqdm(test_loader, desc=\"Predicting\"):\n embeds = embeds.to(device)\n logits = model(embeds)\n probs = torch.sigmoid(logits).cpu().numpy()\n\n all_probs.append(probs)\n all_ids.extend(ids)\n\n all_probs = np.vstack(all_probs) # Shape: (N_test, num_labels)\n all_ids = np.array(all_ids)\n\n # Format predictions for submission\n rows = []\n\n for i, prot_id in enumerate(all_ids):\n probs = all_probs[i]\n\n # Keep only positive probabilities\n idx = np.where(probs > 0)[0]\n if len(idx) == 0:\n continue\n\n # Sort by confidence (descending)\n idx = idx[np.argsort(probs[idx])[::-1]]\n\n # Limit to top 1500 terms per protein \n idx = idx[:max_terms_per_protein]\n\n for j in idx:\n p = probs[j]\n if p <= 0:\n continue\n\n # Format to 3 significant figures \n p_fmt = float(f\"{p:.3g}\")\n if p_fmt <= 0:\n p_fmt = min_prob\n\n rows.append([prot_id, go_terms[j], p_fmt])\n\n # Create submission DataFrame\n sub_df = pd.DataFrame(rows)\n sub_df.to_csv(output_path, sep=\"\\t\", header=False, index=False)\n\n print(f\"\\nāœ… Submission saved to {output_path}\")\n print(f\"Total predictions: {len(sub_df):,}\")\n\n return sub_df\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"GENERATING SUBMISSION\")\nprint(\"=\"*60)\n\nsubmission_df = generate_submission(\n model=model,\n test_loader=test_loader,\n go_terms=top_go,\n device=cfg.device,\n output_path=\"submission.tsv\"\n)\n\n# Display sample predictions\nprint(\"\\nSample predictions:\")\nprint(submission_df.head(10))\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"DONE!\")\nprint(\"=\"*60)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T12:20:03.101508Z","iopub.execute_input":"2026-01-13T12:20:03.101843Z","execution_failed":"2026-01-13T12:21:14.066Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n

\n 9. Ensembling with blast \n

\n","metadata":{}},{"cell_type":"code","source":"# import pandas as pd\n# import numpy as np\n\n\n# submission2 = pd.read_csv(\n# '/kaggle/input/blast-quick-sprof-zero-pred/submission.tsv',\n# sep='\\t',\n# header=None,\n# names=['Id', 'GO term', 'Confidence2'],\n# dtype={'Id': 'category', 'GO term': 'category', 'Confidence2': 'float32'}\n# )\n\n# submission_df.columns = ['Id', 'GO term', 'Confidence']\n# submission_df['Confidence'] = submission_df['Confidence'].astype('float32')\n\n# subs = submission_df.merge(\n# submission2,\n# on=['Id', 'GO term'],\n# how='left'\n# )\n\n# subs['confidence_combined'] = subs['Confidence2'].combine_first(subs['Confidence'])\n\n# subs[['Id', 'GO term', 'confidence_combined']].to_csv(\n# 'submission.tsv',\n# sep='\\t',\n# header=False,\n# index=False\n# )\n","metadata":{"trusted":true,"execution":{"execution_failed":"2026-01-13T12:21:14.067Z"}},"outputs":[],"execution_count":null}]}