{"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":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":13991224,"sourceType":"datasetVersion","datasetId":8917191}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"References:\n- (EDA) https://www.kaggle.com/code/analyticaobscura/cafa-6-decoding-protein-mysteries\n- (ESM-2 embeddings 320 features) https://www.kaggle.com/code/dalloliogm/compute-protein-embeddings-with-esm2-esm-c/notebook\n- (MLP with ESM2) https://www.kaggle.com/code/jwang2025learning/cafa-6-function-prediction-using-prott5?scriptVersionId=282801093\n\n---","metadata":{}},{"cell_type":"code","source":"!pip install biopython > /dev/null","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:07.312435Z","iopub.execute_input":"2025-12-12T14:40:07.312680Z","iopub.status.idle":"2025-12-12T14:40:12.715730Z","shell.execute_reply.started":"2025-12-12T14:40:07.312661Z","shell.execute_reply":"2025-12-12T14:40:12.714941Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 1: Load CAFA6 files\n\n---","metadata":{}},{"cell_type":"code","source":"# CAFA6 file paths\nTRAIN_TERMS = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\nTRAIN_SEQ = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\nTEST_SEQ = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:20.581998Z","iopub.execute_input":"2025-12-12T14:40:20.582318Z","iopub.status.idle":"2025-12-12T14:40:20.586572Z","shell.execute_reply.started":"2025-12-12T14:40:20.582288Z","shell.execute_reply":"2025-12-12T14:40:20.585890Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio import SeqIO\n\n# Dict {entryId, seq}\ntrain_sequences = {rec.id: str(rec.seq) for rec in SeqIO.parse(TRAIN_SEQ, 'fasta')}\ntest_sequences  = {rec.id: str(rec.seq) for rec in SeqIO.parse(TEST_SEQ,  'fasta')}\n\nprint(f'Loaded {len(train_sequences)} train and {len(test_sequences)} test sequences')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:21.778769Z","iopub.execute_input":"2025-12-12T14:40:21.779011Z","iopub.status.idle":"2025-12-12T14:40:25.115857Z","shell.execute_reply.started":"2025-12-12T14:40:21.778994Z","shell.execute_reply":"2025-12-12T14:40:25.114952Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Train dict:\", list(train_sequences.items())[0])\nprint(\"Test dict:\", list(test_sequences.items())[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:26.353679Z","iopub.execute_input":"2025-12-12T14:40:26.353939Z","iopub.status.idle":"2025-12-12T14:40:26.423332Z","shell.execute_reply.started":"2025-12-12T14:40:26.353919Z","shell.execute_reply":"2025-12-12T14:40:26.422748Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ids = [i.split('|')[1] for i in train_sequences.keys()]\ntest_ids = list(test_sequences.keys())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:27.425841Z","iopub.execute_input":"2025-12-12T14:40:27.426079Z","iopub.status.idle":"2025-12-12T14:40:27.447770Z","shell.execute_reply.started":"2025-12-12T14:40:27.426062Z","shell.execute_reply":"2025-12-12T14:40:27.446985Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 2: Feature extraction\n\n---","metadata":{}},{"cell_type":"code","source":"# Embeddings file paths\nESM_EMBEDDINGS = \"/kaggle/input/esm2-t6-8m-ur50d-cafa6\"\nTRAIN_EMBEDDINGS = ESM_EMBEDDINGS + \"/protein_embeddings_train.npy\"\nTEST_EMBEDDINGS = ESM_EMBEDDINGS + \"/protein_embeddings_test.npy\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:37.334095Z","iopub.execute_input":"2025-12-12T14:40:37.334671Z","iopub.status.idle":"2025-12-12T14:40:37.338078Z","shell.execute_reply.started":"2025-12-12T14:40:37.334645Z","shell.execute_reply":"2025-12-12T14:40:37.337502Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\n# Load embeddings\nX_train = np.load(TRAIN_EMBEDDINGS)\nX_test = np.load(TEST_EMBEDDINGS)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:51.931344Z","iopub.execute_input":"2025-12-12T14:40:51.931995Z","iopub.status.idle":"2025-12-12T14:40:55.518564Z","shell.execute_reply.started":"2025-12-12T14:40:51.931971Z","shell.execute_reply":"2025-12-12T14:40:55.517775Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"X_train shape:\", X_train.shape)\nprint(\"X_test shape:\", X_test.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:40:55.519753Z","iopub.execute_input":"2025-12-12T14:40:55.519998Z","iopub.status.idle":"2025-12-12T14:40:55.523804Z","shell.execute_reply.started":"2025-12-12T14:40:55.519972Z","shell.execute_reply":"2025-12-12T14:40:55.523078Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Normalization\nfrom sklearn.preprocessing import StandardScaler\n\nscaler = StandardScaler()\nX_train_norm = scaler.fit_transform(X_train)\nX_test_norm = scaler.transform(X_test)\nprint(\"X_train_norm shape:\", X_train_norm.shape)\nprint(\"X_test_norm shape:\", X_test_norm.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:41:00.804003Z","iopub.execute_input":"2025-12-12T14:41:00.804353Z","iopub.status.idle":"2025-12-12T14:41:02.266880Z","shell.execute_reply.started":"2025-12-12T14:41:00.804328Z","shell.execute_reply":"2025-12-12T14:41:02.266044Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 3: Customized MLP\n\n---","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:41:06.518948Z","iopub.execute_input":"2025-12-12T14:41:06.519386Z","iopub.status.idle":"2025-12-12T14:41:09.856486Z","shell.execute_reply.started":"2025-12-12T14:41:06.519355Z","shell.execute_reply":"2025-12-12T14:41:09.855617Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MLPClassifier(nn.Module):\n    def __init__(self, input_dim, output_dim):\n        super().__init__()\n\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, 2048),\n            nn.ReLU(),\n            nn.BatchNorm1d(2048),\n            nn.Dropout(0.3),\n\n            nn.Linear(2048, 1024),\n            nn.ReLU(),\n            nn.BatchNorm1d(1024),\n            nn.Dropout(0.25),\n\n            nn.Linear(1024, output_dim)  # logits\n        )\n\n    def forward(self, x):\n        return self.net(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:41:09.857683Z","iopub.execute_input":"2025-12-12T14:41:09.858005Z","iopub.status.idle":"2025-12-12T14:41:09.862797Z","shell.execute_reply.started":"2025-12-12T14:41:09.857985Z","shell.execute_reply":"2025-12-12T14:41:09.862130Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 4: Label encoding\n\n---","metadata":{}},{"cell_type":"code","source":"from sklearn.preprocessing import MultiLabelBinarizer\nfrom tqdm import tqdm\nimport pandas as pd\n\ny_trains = dict()\nmlb_dict = dict()\nmodels = dict()\n\ntrain_terms_df = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\n\nfor aspect in ['P', 'C', 'F']:\n    # Filter the train_terms_df based on aspect\n    ont_terms_df = train_terms_df[train_terms_df['aspect'] == aspect]\n\n    # Group the dataFrame based on the EntryID, turn all the GO terms to a list, finally turns this dataFrame to a dict {entryID: terms}\n    protein_terms = ont_terms_df.groupby('EntryID')['term'].apply(list).to_dict()\n\n    # Create a list of labels for this aspect, if an entryID doesn't exist in this aspect, give it a []\n    # This ensures y_train is of shape (82404, ...)\n    labels = [protein_terms.get(entry_id, []) for entry_id in train_ids]\n\n    # Multi-hot encoding, use sparse representation\n    mlb = MultiLabelBinarizer(sparse_output=True)\n    y_train = mlb.fit_transform(labels)\n    y_trains[aspect] = y_train\n    \n    print(f\"y_train shape for {aspect} ontology: {y_train.shape} \\t\\t Number of unique {aspect} terms: {y_train.shape[1]}\")\n\n    # Save to dict\n    mlb_dict[aspect] = mlb\n    model = MLPClassifier(input_dim=X_train.shape[1], output_dim=y_train.shape[1])\n    models[aspect] = model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:42:01.608882Z","iopub.execute_input":"2025-12-12T14:42:01.609183Z","iopub.status.idle":"2025-12-12T14:42:05.974955Z","shell.execute_reply.started":"2025-12-12T14:42:01.609142Z","shell.execute_reply":"2025-12-12T14:42:05.974106Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:42:07.072604Z","iopub.execute_input":"2025-12-12T14:42:07.072888Z","iopub.status.idle":"2025-12-12T14:42:07.079986Z","shell.execute_reply.started":"2025-12-12T14:42:07.072867Z","shell.execute_reply":"2025-12-12T14:42:07.079177Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_trains","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:42:07.365358Z","iopub.execute_input":"2025-12-12T14:42:07.365653Z","iopub.status.idle":"2025-12-12T14:42:07.370591Z","shell.execute_reply.started":"2025-12-12T14:42:07.365630Z","shell.execute_reply":"2025-12-12T14:42:07.369962Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 5: Model training\n\n---","metadata":{}},{"cell_type":"code","source":"# DataLoader\nfrom torch.utils.data import DataLoader, TensorDataset\n\nX_train_tensor = torch.tensor(X_train_norm, dtype=torch.float32)\n\nloaders = {}\n\nfor aspect in ['P', 'C', 'F']:\n    # Convert sparse CSR → dense numpy → float tensor\n    y_dense = y_trains[aspect].toarray().astype('float32')\n    y_tensor = torch.tensor(y_dense, dtype=torch.float32)\n\n    dataset = TensorDataset(X_train_tensor, y_tensor)\n\n    loaders[aspect] = DataLoader(dataset, batch_size=128, shuffle=True)\n\n    print(aspect, \"loader ready:\", y_tensor.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:42:35.160335Z","iopub.execute_input":"2025-12-12T14:42:35.160910Z","iopub.status.idle":"2025-12-12T14:42:46.767413Z","shell.execute_reply.started":"2025-12-12T14:42:35.160886Z","shell.execute_reply":"2025-12-12T14:42:46.766466Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training\nimport torch.optim as optim\n\ncriterion = nn.BCEWithLogitsLoss()\n\ndef train_one_model(model, loader, epochs=5, lr=1e-3, device='cuda'):\n    model = model.to(device)\n    optimizer = optim.Adam(model.parameters(), lr=lr)\n\n    model.train()\n    for ep in range(epochs):\n        total_loss = 0.0\n\n        for X_batch, y_batch in loader:\n            X_batch = X_batch.to(device)\n            y_batch = y_batch.to(device)\n\n            optimizer.zero_grad()\n            logits = model(X_batch)               # shape: (batch, num_labels)\n            loss = criterion(logits, y_batch)\n\n            loss.backward()\n            optimizer.step()\n\n            total_loss += loss.item()\n\n        print(f\"Epoch {ep+1}/{epochs}   Loss = {total_loss/len(loader):.4f}\")\n\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:42:46.768628Z","iopub.execute_input":"2025-12-12T14:42:46.768877Z","iopub.status.idle":"2025-12-12T14:42:46.777321Z","shell.execute_reply.started":"2025-12-12T14:42:46.768860Z","shell.execute_reply":"2025-12-12T14:42:46.776268Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Train 3 models\ntrained_models = dict()\n\nfor aspect in ['P', 'C', 'F']:\n    print(\"\\nTraining\", aspect, \"...\")\n    trained_models[aspect] = train_one_model(\n        model=models[aspect],\n        loader=loaders[aspect],\n        epochs=8,\n        lr=1e-3,\n        device=device\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:42:49.309309Z","iopub.execute_input":"2025-12-12T14:42:49.310055Z","iopub.status.idle":"2025-12-12T14:45:09.207571Z","shell.execute_reply.started":"2025-12-12T14:42:49.310029Z","shell.execute_reply":"2025-12-12T14:45:09.206671Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 6: Inference and Submission\n\n---","metadata":{}},{"cell_type":"code","source":"BATCH_SIZE = 5000  # avoid memory overflow\nsubmission_list = []\n\nfor i in tqdm(range(0, len(test_ids), BATCH_SIZE), desc=\"Predicting on Test Set\"):\n    batch_entry_ids = test_ids[i : i + BATCH_SIZE]\n\n    # Slice features\n    X_batch = X_test_norm[i : i + BATCH_SIZE]\n    X_batch = torch.tensor(X_batch, dtype=torch.float32, device=device)\n\n    # For each ontology aspect (P, F, C)\n    for aspect, model in trained_models.items():\n        model.eval()\n\n        # Forward pass (logits → probabilities)\n        with torch.no_grad():\n            logits = model(X_batch)\n            probs = torch.sigmoid(logits).cpu().numpy()  # ndarray (batch, num_labels)\n\n        mlb = mlb_dict[aspect]  # MultiLabelBinarizer for this aspect\n\n        # Loop over proteins in the batch\n        for j, entry_id in enumerate(batch_entry_ids):\n            prob_vec = probs[j]\n\n            # threshold at 0.02\n            candidate_indices = np.where(prob_vec > 0.02)[0]\n\n            for idx in candidate_indices:\n                submission_list.append(\n                    (entry_id, mlb.classes_[idx], round(prob_vec[idx], 3))\n                )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:45:11.361551Z","iopub.execute_input":"2025-12-12T14:45:11.362181Z","iopub.status.idle":"2025-12-12T14:46:22.878840Z","shell.execute_reply.started":"2025-12-12T14:45:11.362140Z","shell.execute_reply":"2025-12-12T14:46:22.878053Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df = pd.DataFrame(\n    submission_list,\n    columns=['Protein Id', 'GO Term Id', 'Prediction']\n)\n\nsubmission_df = submission_df.sort_values(\n    by=['Protein Id', 'Prediction'],\n    ascending=[True, False]\n)\n\n# Limit 1500 predictions per protein\nfinal_submission_df = (\n    submission_df.groupby('Protein Id')\n    .head(1500)\n    .reset_index(drop=True)\n)\n\nfinal_submission_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:46:22.880064Z","iopub.execute_input":"2025-12-12T14:46:22.880359Z","iopub.status.idle":"2025-12-12T14:46:34.969548Z","shell.execute_reply.started":"2025-12-12T14:46:22.880340Z","shell.execute_reply":"2025-12-12T14:46:34.968718Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\nSubmission file 'submission.tsv' created successfully.\")\nprint(f\"Total predictions in final submission: {len(final_submission_df):,}\")\nprint(\"Submission DataFrame Head:\")\ndisplay(final_submission_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T14:46:34.970379Z","iopub.execute_input":"2025-12-12T14:46:34.970613Z","iopub.status.idle":"2025-12-12T14:46:34.987818Z","shell.execute_reply.started":"2025-12-12T14:46:34.970596Z","shell.execute_reply":"2025-12-12T14:46:34.987264Z"}},"outputs":[],"execution_count":null}]}