{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":116062,"databundleVersionId":14875579},{"sourceType":"datasetVersion","sourceId":5499219,"datasetId":3167603,"databundleVersionId":5573606}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install torch-geometric","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:54.787217Z","iopub.execute_input":"2026-08-16T06:12:54.788389Z","iopub.status.idle":"2026-08-16T06:12:58.345512Z","shell.execute_reply.started":"2026-08-16T06:12:54.788320Z","shell.execute_reply":"2026-08-16T06:12:58.344467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import average_precision_score\nimport optuna\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch_geometric.nn import GATConv\nfrom tqdm.auto import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:58.347588Z","iopub.execute_input":"2026-08-16T06:12:58.347978Z","iopub.status.idle":"2026-08-16T06:12:58.353535Z","shell.execute_reply.started":"2026-08-16T06:12:58.347946Z","shell.execute_reply":"2026-08-16T06:12:58.352749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.manual_seed(42)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:58.354326Z","iopub.execute_input":"2026-08-16T06:12:58.354629Z","iopub.status.idle":"2026-08-16T06:12:58.374425Z","shell.execute_reply.started":"2026-08-16T06:12:58.354603Z","shell.execute_reply":"2026-08-16T06:12:58.373752Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:58.376236Z","iopub.execute_input":"2026-08-16T06:12:58.376604Z","iopub.status.idle":"2026-08-16T06:12:58.384697Z","shell.execute_reply.started":"2026-08-16T06:12:58.376546Z","shell.execute_reply":"2026-08-16T06:12:58.384014Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"NUM_TOP_TERMS = 1500","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:58.385831Z","iopub.execute_input":"2026-08-16T06:12:58.386204Z","iopub.status.idle":"2026-08-16T06:12:58.396145Z","shell.execute_reply.started":"2026-08-16T06:12:58.386176Z","shell.execute_reply":"2026-08-16T06:12:58.395256Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_embeds = np.load(\"/kaggle/input/datasets/sergeifironov/t5embeds/train_embeds.npy\")\ntrain_ids = np.load(\"/kaggle/input/datasets/sergeifironov/t5embeds/train_ids.npy\")\n\ntrain_terms = pd.read_csv(\"/kaggle/input/competitions/cafa-6-protein-function-prediction/Train/train_terms.tsv\", sep= \"\\t\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:58.397217Z","iopub.execute_input":"2026-08-16T06:12:58.397478Z","iopub.status.idle":"2026-08-16T06:12:58.971816Z","shell.execute_reply.started":"2026-08-16T06:12:58.397453Z","shell.execute_reply":"2026-08-16T06:12:58.971203Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"term_counts = train_terms[\"term\"].value_counts()\ntop_terms = term_counts.head(NUM_TOP_TERMS).index.tolist()\n\nterm_to_idx = {term: idx for idx, term in enumerate(top_terms)}\n\ntrain_terms_filtered = train_terms[train_terms[\"term\"].isin(top_terms)]\n\nprotein_to_terms = train_terms_filtered.groupby(\"EntryID\")[\"term\"].apply(list).to_dict()\n\ntargets = np.zeros((len(train_ids), NUM_TOP_TERMS), dtype= np.float32)\n\nfor i, prot_id in enumerate(tqdm(train_ids, desc= \"Building targets\")):\n    if prot_id in protein_to_terms:\n        for term in protein_to_terms[prot_id]:\n            targets[i, term_to_idx[term]] = 1.0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:12:58.972803Z","iopub.execute_input":"2026-08-16T06:12:58.973127Z","iopub.status.idle":"2026-08-16T06:13:01.067444Z","shell.execute_reply.started":"2026-08-16T06:12:58.973097Z","shell.execute_reply":"2026-08-16T06:13:01.066607Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"obo_path = \"/kaggle/input/competitions/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n\nedges = []\nwith open(obo_path, 'r') as f:\n    current_term = None\n    for line in f:\n        line = line.strip()\n        if line.startswith(\"id: GO:\"):\n            current_term = line.split(\"id: \")[1]\n        elif line.startswith(\"is_a: GO:\") and current_term in term_to_idx:\n            parent_term = line.split(\"is_a: \")[1].split(\" !\")[0]\n            if parent_term in term_to_idx:\n                parent_idx = term_to_idx[parent_term]\n                child_idx = term_to_idx[current_term]\n                edges.append((parent_idx, child_idx))\n                edges.append((child_idx, parent_idx))\n\n\nedge_index = torch.tensor(edges, dtype= torch.long).t().contiguous().to(device)\nprint(f\"Graph built with {NUM_TOP_TERMS} nodes and {edge_index.shape[1]} edges.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:13:01.068756Z","iopub.execute_input":"2026-08-16T06:13:01.069260Z","iopub.status.idle":"2026-08-16T06:13:01.300901Z","shell.execute_reply.started":"2026-08-16T06:13:01.069207Z","shell.execute_reply":"2026-08-16T06:13:01.300185Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_val_x, test_x, train_val_y, test_y = train_test_split(\n    train_embeds, targets, test_size= 0.10, random_state= 42\n)\n\ntrain_x, val_x, train_y, val_y = train_test_split(\n    train_val_x, train_val_y, test_size= 1/9, random_state= 42\n)\n\nclass ProteinDataset(Dataset):\n    def __init__(self, embeddings, targets):\n        self.embeddings = torch.tensor(embeddings, dtype= torch.float32)\n        self.targets = torch.tensor(targets, dtype= torch.float32)\n        \n    def __len__(self):\n        return len(self.embeddings)\n    \n    def __getitem__(self, idx):\n        return self.embeddings[idx], self.targets[idx]\n\ntrain_dataset = ProteinDataset(train_x, train_y)\nval_dataset = ProteinDataset(val_x, val_y)\ntest_dataset = ProteinDataset(test_x, test_y)\n\ntrain_loader = DataLoader(train_dataset, batch_size= 256, shuffle= True)\nval_loader = DataLoader(val_dataset, batch_size= 256, shuffle= False)\ntest_loader = DataLoader(test_dataset, batch_size= 256, shuffle= False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:13:01.301938Z","iopub.execute_input":"2026-08-16T06:13:01.302249Z","iopub.status.idle":"2026-08-16T06:13:03.228996Z","shell.execute_reply.started":"2026-08-16T06:13:01.302221Z","shell.execute_reply":"2026-08-16T06:13:03.228347Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProtGO_GAT(nn.Module):\n    def __init__(self, num_protein_features= 1024, num_go_terms= NUM_TOP_TERMS, hidden_dim= 256, dropout_rate= 0.3):\n        super(ProtGO_GAT, self).__init__()\n        \n        self.protein_mlp = nn.Sequential(\n            nn.Linear(num_protein_features, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(dropout_rate),\n            nn.Linear(512, hidden_dim)\n        )\n        \n        self.go_node_embeds = nn.Parameter(torch.Tensor(num_go_terms, hidden_dim))\n        nn.init.xavier_uniform_(self.go_node_embeds)\n        \n        self.gat1 = GATConv(hidden_dim, hidden_dim // 4, heads= 4, dropout= dropout_rate)\n        self.gat2 = GATConv(hidden_dim, hidden_dim // 4, heads= 4, dropout= dropout_rate)\n        \n    def forward(self, protein_features, edge_index):\n        p_emb = self.protein_mlp(protein_features)\n\n        g_initial = self.go_node_embeds\n        \n        g_update = self.gat1(g_initial, edge_index)\n        g_update = torch.relu(g_update)\n        g_update = self.gat2(g_update,edge_index)\n\n        g_emb = g_initial + g_update\n\n        logits = torch.matmul(p_emb, g_emb.t())\n        return logits","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:13:03.230850Z","iopub.execute_input":"2026-08-16T06:13:03.231058Z","iopub.status.idle":"2026-08-16T06:13:03.237596Z","shell.execute_reply.started":"2026-08-16T06:13:03.231034Z","shell.execute_reply":"2026-08-16T06:13:03.236815Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def objective(trial):\n    lr = trial.suggest_float(\"lr\", 1e-4, 5e-3, log= True)\n    hidden_dim = trial.suggest_categorical(\"hidden_dim\", [128, 256, 512])\n    dropout_rate = trial.suggest_float(\"dropout_rate\", 0.1, 0.4)\n\n    pos_weight_val = trial.suggest_float(\"pos_weight\", 5.0, 50.0) \n\n    model = ProtGO_GAT(hidden_dim= hidden_dim, dropout_rate= dropout_rate).to(device)\n    optimiser = torch.optim.Adam(model.parameters(), lr= lr, weight_decay= 1e-5)\n\n    pos_weight = torch.tensor([pos_weight_val]).to(device)\n    criterion = nn.BCEWithLogitsLoss(pos_weight= pos_weight)\n\n    epochs = 5\n    for epoch in range(epochs):\n        model.train()\n        for batch_embeds, batch_targets in train_loader:\n            batch_embeds, batch_targets = batch_embeds.to(device), batch_targets.to(device)\n            \n            optimiser.zero_grad()\n            logits = model(batch_embeds, edge_index)\n            loss = criterion(logits, batch_targets)\n            loss.backward()\n            optimiser.step()\n\n    model.eval()\n    val_preds, val_trues = [], []\n    with torch.no_grad():\n        for batch_embeds, batch_targets in val_loader:\n            batch_embeds = batch_embeds.to(device)\n            probs = torch.sigmoid(model(batch_embeds, edge_index))\n            val_preds.append(probs.cpu().numpy())\n            val_trues.append(batch_targets.numpy())\n            \n    val_preds = np.vstack(val_preds)\n    val_trues = np.vstack(val_trues)\n\n    macro_auprc = average_precision_score(val_trues, val_preds, average= \"macro\")\n    return macro_auprc\n\nstudy = optuna.create_study(direction= \"maximize\", sampler=optuna.samplers.TPESampler(seed=42))\nstudy.optimize(objective, n_trials= 10)\n\nprint(f\"Best Hyperparameters Found: {study.best_params}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:13:03.238655Z","iopub.execute_input":"2026-08-16T06:13:03.239286Z","iopub.status.idle":"2026-08-16T06:16:49.727026Z","shell.execute_reply.started":"2026-08-16T06:13:03.239258Z","shell.execute_reply":"2026-08-16T06:16:49.726187Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"final_x = np.concatenate([train_x, val_x], axis= 0)\nfinal_y = np.concatenate([train_y, val_y], axis= 0)\n\nfinal_dataset = ProteinDataset(final_x, final_y)\n\nfinal_loader = DataLoader(final_dataset, batch_size= 256, shuffle= True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:16:49.728173Z","iopub.execute_input":"2026-08-16T06:16:49.728995Z","iopub.status.idle":"2026-08-16T06:16:50.950982Z","shell.execute_reply.started":"2026-08-16T06:16:49.728960Z","shell.execute_reply":"2026-08-16T06:16:50.950084Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"best_params = study.best_params\n\nNUM_MODELS = 3\nepochs = 25\nensemble_models = []\n\n\nfor i in range(NUM_MODELS):\n    print(f\"\\n--- Training Model {i + 1}/{NUM_MODELS} ---\")\n\n    torch.manual_seed(42 + i)\n\n    model = ProtGO_GAT(hidden_dim= best_params[\"hidden_dim\"], dropout_rate= best_params[\"dropout_rate\"]).to(device)\n    optimiser = torch.optim.Adam(model.parameters(), lr= best_params[\"lr\"], weight_decay= 1e-5)\n\n    pos_weight = torch.tensor([best_params[\"pos_weight\"]]).to(device)\n    criterion = nn.BCEWithLogitsLoss(pos_weight= pos_weight)\n\n    for epoch in range(epochs):\n        model.train()\n        for batch_embeds, batch_targets in tqdm(final_loader, desc= f\"Epoch {epoch + 1}/{epochs}\"):\n            batch_embeds, batch_targets = batch_embeds.to(device), batch_targets.to(device)\n            optimiser.zero_grad()\n            logits = model(batch_embeds, edge_index)\n            loss = criterion(logits, batch_targets)\n            loss.backward()\n            optimiser.step()\n\n    model.eval()\n    ensemble_models.append(model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:16:50.952022Z","iopub.execute_input":"2026-08-16T06:16:50.952425Z","iopub.status.idle":"2026-08-16T06:22:14.670552Z","shell.execute_reply.started":"2026-08-16T06:16:50.952376Z","shell.execute_reply":"2026-08-16T06:22:14.669572Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_preds = []\ntest_trues = []\n\nwith torch.no_grad():\n    for batch_embeds, batch_targets in tqdm(test_loader, desc= \"Evaluating Ensemble\"):\n        batch_embeds = batch_embeds.to(device)\n\n        model_probs = []\n        for m in ensemble_models:\n            logits = m(batch_embeds, edge_index)\n            model_probs.append(torch.sigmoid(logits))\n\n        avg_probs = torch.stack(model_probs).mean(dim= 0)\n        \n        test_preds.append(avg_probs.cpu().numpy())\n        test_trues.append(batch_targets.numpy())\n\ntest_preds = np.vstack(test_preds)\ntest_trues = np.vstack(test_trues)\n\nmacro_auprc = average_precision_score(test_trues, test_preds, average= \"macro\")\nmicro_auprc = average_precision_score(test_trues, test_preds, average= \"micro\")\n\nprint(f\"\\n--- Final Ensemble Test Results ---\")\nprint(f\"Macro AUPRC: {macro_auprc:.4f}\")\nprint(f\"Micro AUPRC: {micro_auprc:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-08-16T06:22:14.671751Z","iopub.execute_input":"2026-08-16T06:22:14.672100Z","iopub.status.idle":"2026-08-16T06:22:25.600788Z","shell.execute_reply.started":"2026-08-16T06:22:14.672039Z","shell.execute_reply":"2026-08-16T06:22:25.600016Z"},"scrolled":true},"outputs":[],"execution_count":null}]}