{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"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"},{"sourceId":14403269,"sourceType":"datasetVersion","datasetId":9198987}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%pip install biopython","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:45:32.355546Z","iopub.execute_input":"2026-01-05T22:45:32.356323Z","iopub.status.idle":"2026-01-05T22:45:37.822798Z","shell.execute_reply.started":"2026-01-05T22:45:32.356297Z","shell.execute_reply":"2026-01-05T22:45:37.821906Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom Bio import SeqIO\nimport unicodedata\nimport re\nfrom joblib import Parallel, delayed\nfrom time import time\nfrom collections import defaultdict\nfrom sklearn.neighbors import NearestNeighbors\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport random\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:45:45.102185Z","iopub.execute_input":"2026-01-05T22:45:45.102665Z","iopub.status.idle":"2026-01-05T22:45:49.473804Z","shell.execute_reply.started":"2026-01-05T22:45:45.102616Z","shell.execute_reply":"2026-01-05T22:45:49.473155Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 1. Data Preparation","metadata":{}},{"cell_type":"code","source":"# load the train sequences\ntrain_sequence_file = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\ntotal = sum(1 for line in open(train_sequence_file, \"r\", encoding=\"utf-8\") if line.startswith(\">\"))\nprint(f\"Total number of train records: {total}\")\n\ntrain_sequence_df = []\nfor record in tqdm(SeqIO.parse(train_sequence_file, \"fasta\"), total=total, desc=\"Parsing FASTA\", unit=\"seq\"):\n    r = [\n        record.id, \n        record.id.split('|')[1],\n        record.description, \n        str(record.seq)\n    ]\n    train_sequence_df.append(r)\n    \n# convert to pandas data frame\ntrain_sequence_df = pd.DataFrame(train_sequence_df, columns=['full_prot_id', 'prot_id', 'description', 'sequence'])\ntrain_sequence_df_backup = train_sequence_df.copy()\ntrain_sequence_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:45:53.342446Z","iopub.execute_input":"2026-01-05T22:45:53.343243Z","iopub.status.idle":"2026-01-05T22:45:54.893331Z","shell.execute_reply.started":"2026-01-05T22:45:53.343217Z","shell.execute_reply":"2026-01-05T22:45:54.892672Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load train-terms\ntrain_term_df = pd.read_csv(\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n    sep=\"\\t\"\n)\ntrain_term_df = train_term_df.rename(columns={'EntryID': 'prot_id'})\ntrain_term_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:00.230753Z","iopub.execute_input":"2026-01-05T22:46:00.231402Z","iopub.status.idle":"2026-01-05T22:46:00.579638Z","shell.execute_reply.started":"2026-01-05T22:46:00.231378Z","shell.execute_reply":"2026-01-05T22:46:00.578883Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# merge to get labels\ntrain_sequence_df = train_sequence_df.merge(\n    train_term_df,\n    on='prot_id'\n)\ntrain_sequence_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:04.585709Z","iopub.execute_input":"2026-01-05T22:46:04.586285Z","iopub.status.idle":"2026-01-05T22:46:04.793478Z","shell.execute_reply.started":"2026-01-05T22:46:04.586260Z","shell.execute_reply":"2026-01-05T22:46:04.792912Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# load test data\ntest_superset_file = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\n# Count sequences quickly (without parsing)\ntotal = sum(1 for line in open(test_superset_file, \"r\", encoding=\"utf-8\") if line.startswith(\">\"))\nprint(f\"Total number of records: {total}\")\n\n# Load all sequences from a FASTA file\ntest_df = []\nfor record in tqdm(SeqIO.parse(test_superset_file, \"fasta\"), total=total, desc=\"Parsing FASTA\", unit=\"seq\"):\n    r = [record.id, record.description, str(record.seq)]\n    test_df.append(r)\n\ntest_df = pd.DataFrame(test_df, columns=['prot_id', 'description', 'sequence'])\nprint(f\"No. proteins: {test_df.prot_id.nunique()}\")\nprint(f\"No. sequences: {test_df.sequence.nunique()}\")\ntest_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:18.123168Z","iopub.execute_input":"2026-01-05T22:46:18.123709Z","iopub.status.idle":"2026-01-05T22:46:21.400870Z","shell.execute_reply.started":"2026-01-05T22:46:18.123664Z","shell.execute_reply":"2026-01-05T22:46:21.400321Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 2. Feature","metadata":{}},{"cell_type":"code","source":"# load features\nX_train = np.load(\"/kaggle/input/protein-cafa6-embedding/esm2_t30_150M_UR50D/train.npy\")\nX_test = np.load(\"/kaggle/input/protein-cafa6-embedding/esm2_t30_150M_UR50D/test.npy\")\nprint(f\"X_train shape: {X_train.shape}, X_test shape: {X_test.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:38.890882Z","iopub.execute_input":"2026-01-05T22:46:38.891188Z","iopub.status.idle":"2026-01-05T22:46:42.880195Z","shell.execute_reply.started":"2026-01-05T22:46:38.891166Z","shell.execute_reply":"2026-01-05T22:46:42.879343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train = train_sequence_df.drop_duplicates(subset=['prot_id']).reset_index(drop=True)\ntrain","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:46.020384Z","iopub.execute_input":"2026-01-05T22:46:46.020992Z","iopub.status.idle":"2026-01-05T22:46:46.090805Z","shell.execute_reply.started":"2026-01-05T22:46:46.020966Z","shell.execute_reply":"2026-01-05T22:46:46.090018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# prepare labels\nLABEL_NUM_LIMIT = 4096\ny_train = train_sequence_df[[\"prot_id\", \"term\"]].dropna().drop_duplicates()\nprint(f\"No. rows y_train: {len(y_train)}\")\nterm_frequency = y_train.term.value_counts()\nterm_frequency = term_frequency.iloc[:LABEL_NUM_LIMIT]\nnum_classes = len(term_frequency)\nvocab = term_frequency.index.tolist()\nterm2index = {t: int(i) for i, t in enumerate(vocab)}\ny_train['term_indx'] = y_train['term'].map(term2index)\n# y_train['term_indx'] = y_train['term_indx'].astype(int)\nprint(f\"No. rows without a label: {y_train.term_indx.isna().sum()}\")\ny_train = y_train[~y_train.term_indx.isna()]\ny_train['term_indx'] = y_train['term_indx'].astype(int)\nprint(f\"No. rows y_train top k left: {len(y_train)}\")\ny_train = y_train.groupby(\"prot_id\")['term_indx'].apply(list).reset_index()\n\n# filter X_train\nprot2indx = dict(zip(train.prot_id, train.index))\nkeep_x_indices = y_train.prot_id.map(prot2indx)\nX_train = X_train[keep_x_indices]\n\n# convert to list of lists\ny_train = y_train.term_indx.tolist()\nassert len(y_train) == X_train.shape[0], \"missing labels!\"\nprint(f\"Total number of classes: {num_classes}\")\nprint(f\"Total number of labeled examples: {len(y_train)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:48.104289Z","iopub.execute_input":"2026-01-05T22:46:48.104580Z","iopub.status.idle":"2026-01-05T22:46:49.874553Z","shell.execute_reply.started":"2026-01-05T22:46:48.104558Z","shell.execute_reply":"2026-01-05T22:46:49.873833Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Vocab size: {len(term2index)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:53.231628Z","iopub.execute_input":"2026-01-05T22:46:53.232450Z","iopub.status.idle":"2026-01-05T22:46:53.236405Z","shell.execute_reply.started":"2026-01-05T22:46:53.232422Z","shell.execute_reply":"2026-01-05T22:46:53.235563Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# data splitting\nfrom sklearn.model_selection import train_test_split\n\nX_train, X_val, y_train, y_val = train_test_split(\n    X_train, y_train, test_size=0.1, random_state=42)\nprint(f\"No. train examples: {len(X_train)}\")\nprint(f\"No. val examples: {len(X_val)}\")\nassert len(X_train) == len(y_train), \"train mismatch!\"\nassert len(X_val) == len(y_val), \"train mismatch!\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:55.997546Z","iopub.execute_input":"2026-01-05T22:46:55.998081Z","iopub.status.idle":"2026-01-05T22:46:56.094541Z","shell.execute_reply.started":"2026-01-05T22:46:55.998057Z","shell.execute_reply":"2026-01-05T22:46:56.093713Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### 3. Neural Net","metadata":{}},{"cell_type":"code","source":"class MLPClassifier(nn.Module):\n    def __init__(self, d, h, C):\n        super().__init__()\n        self.encoder = nn.Sequential(\n            nn.Linear(d, h),\n            nn.ReLU(),\n            nn.Dropout(0.1),\n            nn.Linear(h, h),\n            nn.ReLU(),\n        )\n        self.out = nn.Linear(h, C)  # logits for each label\n\n    def forward(self, x):\n        z = self.encoder(x)\n        logits = self.out(z)  # [B, C]\n        return logits, z\n\n\ndef sampled_bce_with_logits_batched(\n    logits,\n    pos_indices_per_sample,\n    C,\n    num_negatives=200,\n    device=None,\n    generator=None,\n):\n    \"\"\"\n    Batched negative sampling using row-wise top-k.\n    \"\"\"\n    B = logits.size(0)\n    device = device or logits.device\n\n    # Build [B, C] positive mask\n    pos_mask = torch.zeros((B, C), dtype=torch.bool, device=device)\n    for i, pos_i in enumerate(pos_indices_per_sample):\n        if len(pos_i) > 0:\n            pos_mask[i, torch.tensor(pos_i, dtype=torch.long, device=device)] = True\n\n    # Random scores for all classes; suppress positives\n    rand_scores = torch.rand((B, C), device=device, generator=generator)\n    rand_scores = rand_scores.masked_fill(pos_mask, -1e9)\n\n    # Row-wise top-k (same k for all rows). If some rows have fewer availables,\n    # we will filter positives out later.\n    k = min(num_negatives, int((~pos_mask).sum(dim=1).min().item()))\n    neg_candidates = torch.topk(rand_scores, k=k, dim=1, sorted=False).indices  # [B, k]\n\n    loss_list = []\n    for i in range(B):\n        pos_i = torch.tensor(pos_indices_per_sample[i], dtype=torch.long, device=device)\n        # filter any accidental positives (shouldn't occur because of -1e9 and k <= min available)\n        neg_i = neg_candidates[i]\n        # Combine\n        indices = torch.cat([pos_i, neg_i])\n        targets = torch.zeros(indices.size(0), device=device)\n        targets[:pos_i.numel()] = 1.0\n        selected_logits = logits[i].index_select(0, indices)\n        loss_i = F.binary_cross_entropy_with_logits(selected_logits, targets, reduction='mean')\n        loss_list.append(loss_i)\n\n    return torch.stack(loss_list).mean()\n\n\n\nimport torch\n\ndef evaluate(model, X, y, num_classes, num_negatives, device='cuda', threshold=0.5):\n    \"\"\"Evaluate a multi-label classifier that outputs logits of shape [B, C].\n    Args:\n        model: a PyTorch neural network. Forward returns (logits, *), where logits is [B, C].\n        X: numpy array of shape [m, d]\n        y: list of lists (length m). Each inner list contains positive class indices for that sample.\n        num_classes: total number of classes (C)\n        num_negatives: number of negatives to sample in the loss (for sampled_bce_with_logits_batched)\n        device: 'cuda' or 'cpu'\n\n    Returns:\n        dict: {\"loss\": mean_loss, \"macro_f1\": macro_f1, \"micro_f1\": micro_f1}\n    \"\"\"\n    model.eval()\n    dataloader = create_sparse_dataloader(X, y, batch_size=512)\n\n    total_loss = 0.0\n    num_batches = 0\n\n    # Accumulators for F1 (kept on device for efficiency)\n    per_class_tp = torch.zeros(num_classes, dtype=torch.long, device=device)\n    per_class_fp = torch.zeros(num_classes, dtype=torch.long, device=device)\n    per_class_fn = torch.zeros(num_classes, dtype=torch.long, device=device)\n\n    with torch.no_grad():\n        for X_batch, pos_indices_list in dataloader:\n            num_batches += 1\n\n            # Move inputs\n            X_batch = X_batch.to(device, non_blocking=True)\n\n            # Forward + sampled loss (mixed precision if CUDA)\n            with torch.cuda.amp.autocast(enabled=(device == 'cuda')):\n                logits, _ = model(X_batch)  # [B, C]\n                loss = sampled_bce_with_logits_batched(\n                    logits,\n                    pos_indices_list,\n                    C=num_classes,\n                    num_negatives=num_negatives,\n                    device=device\n                )\n            total_loss += loss.item()\n\n            # Predictions\n            probs = torch.sigmoid(logits)        # [B, C]\n            preds = (probs >= torch.from_numpy(threshold).to(device)) if isinstance(threshold, np.ndarray) else  (probs >= threshold)          # bool [B, C]\n\n            # Build dense multi-hot targets from sparse index lists\n            targets = torch.zeros_like(preds, dtype=torch.bool)  # [B, C]\n            for i, idxs in enumerate(pos_indices_list):\n                if idxs:  # skip empty label sets\n                    idxs_t = torch.as_tensor(idxs, device=device)\n                    targets[i, idxs_t] = True\n\n            # Per-class confusion components over the batch\n            per_class_tp += (preds & targets).sum(dim=0)\n            per_class_fp += (preds & ~targets).sum(dim=0)\n            per_class_fn += (~preds & targets).sum(dim=0)\n\n    # Mean loss across batches\n    mean_loss = total_loss / max(1, num_batches)\n\n    # Micro-F1: aggregate TP/FP/FN across all classes\n    tp_sum = per_class_tp.sum().item()\n    fp_sum = per_class_fp.sum().item()\n    fn_sum = per_class_fn.sum().item()\n    micro_denom = (2 * tp_sum + fp_sum + fn_sum)\n    micro_f1 = (2 * tp_sum / micro_denom) if micro_denom > 0 else 0.0\n\n    # Macro-F1: average per-class F1\n    # F1_c = 2*TP_c / (2*TP_c + FP_c + FN_c); set F1_c=0 when denominator is 0\n    per_class_denom = (2 * per_class_tp + per_class_fp + per_class_fn).clamp(min=1)\n    macro_f1 = (2 * per_class_tp.float() / per_class_denom.float()).mean().item()\n\n    return {\n        \"loss\": float(mean_loss),\n        \"macro_f1\": float(macro_f1),\n        \"micro_f1\": float(micro_f1),\n    }   \n\n\n# ---- Training loop sketch ----\ndef train(\n    model,\n    dataloader,                    # yields (X_batch [B,d], pos_indices_list [List[List[int]]])\n    C,\n    num_epochs=10,\n    lr=1e-3,\n    num_negatives=200,\n    device='cuda',\n    patience=10,\n    save_path = \"best_multilabel.pt\"\n):\n    model.to(device)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)\n    scaler = torch.cuda.amp.GradScaler(enabled=(device=='cuda'))\n\n    model.train()\n    best_val_f1 = -1.0\n    epochs_without_improvement = 0\n    for epoch in range(num_epochs):\n        total_loss = 0.0\n        for X_batch, pos_indices_list in dataloader:\n            X_batch = X_batch.to(device, non_blocking=True)\n\n            with torch.amp.autocast('cuda', enabled=(device == 'cuda')):\n                logits, _ = model(X_batch)\n                loss = sampled_bce_with_logits_batched(\n                    logits,\n                    pos_indices_list,\n                    C=C,\n                    num_negatives=num_negatives,\n                    device=device\n                )\n\n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad(set_to_none=True)\n\n            total_loss += loss.item()\n        \n        # calculating performance on the validation set\n        val_performance = evaluate(model, X_val, y_val, C, num_negatives)\n        print(f\"\"\"Epoch {epoch+1}: train_loss={total_loss/len(dataloader):.4f}, val_loss={val_performance['loss']:.4f}, val_macro_f1={val_performance['macro_f1']:.4f}, val_micro_f1={val_performance['micro_f1']:.4f}\"\"\")\n        \n        # Early stopping on tuned macro-F1\n        macro_f1 = val_performance['macro_f1']\n        if macro_f1 > best_val_f1 + 1e-6:\n            best_val_f1 = macro_f1\n            epochs_without_improvement = 0\n            torch.save({\n                \"model_state\": model.state_dict()\n            }, save_path)\n        else:\n            epochs_without_improvement += 1\n            if epochs_without_improvement >= patience:\n                print(f\"Early stopping at epoch {epoch} (no improvement for {patience} epochs).\")\n                break\n    print(f\"Best best_val_f1: {best_val_f1}\")\n\n# ---- Inference: get top-k labels ----\ndef predict_topk(model, X_batch, k=10, device='cuda'):\n    model.eval()\n    with torch.no_grad(), torch.cuda.amp.autocast(enabled=(device=='cuda')):\n        logits, _ = model(X_batch.to(device))\n        probs = torch.sigmoid(logits)  # [B, C]\n        # Top-k per sample\n        topk_probs, topk_idx = torch.topk(probs, k, dim=1)\n    return topk_idx.cpu(), topk_probs.cpu()\n\n\ndef predict_all(model, X, device='cuda', batch_size=512):\n    model.eval()\n    probs = []\n    n_test = X.shape[0]\n    for i in tqdm(range(0, n_test, batch_size)):\n        X_batch = torch.Tensor(X[i: i+batch_size])\n        with torch.no_grad(), torch.cuda.amp.autocast(enabled=(device=='cuda')):\n            logits, _ = model(X_batch.to(device))\n            prob_batch = torch.sigmoid(logits).cpu()  # [B, C]\n            probs.append(prob_batch)\n    return np.concatenate(probs, axis=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:46:58.519624Z","iopub.execute_input":"2026-01-05T22:46:58.520225Z","iopub.status.idle":"2026-01-05T22:46:58.549315Z","shell.execute_reply.started":"2026-01-05T22:46:58.520200Z","shell.execute_reply":"2026-01-05T22:46:58.548464Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model init\nmodel = MLPClassifier(d=X_train.shape[1], h=1024, C=num_classes)\nmodel","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:47:06.945951Z","iopub.execute_input":"2026-01-05T22:47:06.946273Z","iopub.status.idle":"2026-01-05T22:47:07.015751Z","shell.execute_reply.started":"2026-01-05T22:47:06.946250Z","shell.execute_reply":"2026-01-05T22:47:07.015094Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# create data loader\nclass SparseMultiLabelDataset(Dataset):\n    def __init__(self, X, pos_indices_per_sample, dtype=torch.float32):\n        \"\"\"\n        X: np.ndarray or torch.Tensor, shape [m, d]\n        pos_indices_per_sample: List[List[int]], length m\n            Each entry is a list of positive label indices for that sample.\n        \"\"\"\n        self.X = torch.as_tensor(X, dtype=dtype)\n        self.pos_indices = pos_indices_per_sample\n        assert self.X.shape[0] == len(self.pos_indices), \"X and labels length mismatch\"\n\n    def __len__(self):\n        return self.X.shape[0]\n\n    def __getitem__(self, idx):\n        # Return features and the list of positive label indices\n        return self.X[idx], self.pos_indices[idx]\n\n# Custom collate to keep lists intact\ndef sparse_collate(batch):\n    \"\"\"\n    batch: list of tuples (X_i, pos_list_i)\n    Returns:\n        X_batch: torch.Tensor [B, d]\n        pos_lists: List[List[int]] length B\n    \"\"\"\n    Xs, pos_lists = zip(*batch)\n    X_batch = torch.stack(Xs, dim=0)\n    return X_batch, list(pos_lists)\n\ndef create_sparse_dataloader(X_train, pos_indices_per_sample, batch_size=256, shuffle=True, num_workers=4):\n    dataset = SparseMultiLabelDataset(X_train, pos_indices_per_sample)\n    loader = DataLoader(\n        dataset,\n        batch_size=batch_size,\n        shuffle=shuffle,\n        num_workers=num_workers,\n        pin_memory=True,\n        persistent_workers=(num_workers > 0),\n        drop_last=False,\n        collate_fn=sparse_collate     # important: keep sparse label lists\n    )\n    return loader\n\n\nloader = create_sparse_dataloader(X_train, y_train, batch_size=512)\nloader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-05T22:47:10.643276Z","iopub.execute_input":"2026-01-05T22:47:10.643578Z","iopub.status.idle":"2026-01-05T22:47:10.659483Z","shell.execute_reply.started":"2026-01-05T22:47:10.643556Z","shell.execute_reply":"2026-01-05T22:47:10.658791Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# training\nt_start = time()\ntrain(\n    model=model,\n    dataloader=loader,                    # yields (X_batch [B,d], pos_indices_list [List[List[int]]])\n    C=num_classes,\n    num_epochs=200,\n    lr=1e-4,\n    num_negatives=500,\n    device='cuda'\n)\nprint(f\"Training finished in {time() - t_start} seconds!\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# reload the best model\nmodel_path = \"best_multilabel.pt\"\nckpt = torch.load(model_path, map_location='cpu')\nmodel = MLPClassifier(d=X_train.shape[1], h=1024, C=num_classes).to('cuda').eval()\nmodel.load_state_dict(ckpt[\"model_state\"])","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# fine-tune thresholds\nval_probs = predict_all(model, X_val)\nprint(f\"val_probs shape: {val_probs.shape}\")\n\ndef to_full_label(y, C):\n    Y = np.zeros((len(y), C))\n    for i, r in tqdm(enumerate(y_val)):\n        Y[i, r] = 1\n    return Y\n\ny_val_full = to_full_label(y_val, num_classes)\nprint(f\"y_val_full shape: {y_val_full.shape}\")\n\ndef find_best_thresholds_numpy(y_true: np.ndarray,\n                               y_prob: np.ndarray,\n                               grid: np.ndarray | None = None,\n                               eps: float = 1e-8) -> np.ndarray:\n    \"\"\"\n    Find per-class thresholds maximizing F1 on validation set (NumPy version).\n\n    Parameters\n    ----------\n    y_true : np.ndarray\n        Shape (N, C), binary ground-truth {0,1}.\n    y_prob : np.ndarray\n        Shape (N, C), sigmoid probabilities in [0,1].\n    grid : np.ndarray | None\n        1D array of thresholds to try. If None, uses np.linspace(0.05, 0.95, 19).\n    eps : float\n        Small constant for numerical stability.\n\n    Returns\n    -------\n    np.ndarray\n        Shape (C,), best threshold for each class.\n    \"\"\"\n    if grid is None:\n        grid = np.linspace(0.05, 0.95, num=19)\n\n    # Ensure proper dtypes\n    y_prob = y_prob.astype(np.float64)\n    grid = grid.astype(np.float64)\n\n    N, C = y_true.shape\n    T = grid.shape[0]\n\n    # Broadcast predictions for all thresholds: shape (T, N, C)\n    pred = (y_prob[None, :, :] >= grid[:, None, None])\n\n    # Precompute positives per class\n    positives = y_true.sum(axis=0)  # shape (C,)\n\n    # Compute TP, FP, FN for all (threshold, class) pairs\n    tp = (pred & (y_true[None, :, :] == 1)).sum(axis=1)         # (T, C)\n    pp = pred.sum(axis=1)                                       # (T, C)\n    fp = pp - tp                                                # (T, C)\n    fn = positives[None, :] - tp                                # (T, C)\n\n    precision = tp / (tp + fp + eps)                            # (T, C)\n    recall = tp / (tp + fn + eps)                               # (T, C)\n    f1 = 2.0 * precision * recall / (precision + recall + eps)  # (T, C)\n\n    # Argmax over thresholds per class\n    best_idx = np.argmax(f1, axis=0)                            # (C,)\n    best_t = grid[best_idx]                                     # (C,)\n\n    return best_t\n\n\nprint(\"Detecting best thresholds...\")\nbest_thresholds = find_best_thresholds_numpy(y_val_full, val_probs)\nbest_thresholds","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# re-calculate the performance using the detected thresholds\nval_performance_tuned = evaluate(model, X_val, y_val, num_classes, num_negatives=500, threshold=best_thresholds)\nval_performance_tuned","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# making prediction using topk\nN_NEIGHBORS = 50\nn_test = X_test.shape[0]\nBATCH_SIZE = 512\ntopk = []\nprob = []\nfor i in tqdm(range(0, n_test, BATCH_SIZE)):\n    X_batch = torch.Tensor(X_test[i: i+BATCH_SIZE])\n    topk_batch, prob_batch = predict_topk(\n        model, \n        X_batch, \n        k=N_NEIGHBORS, \n        device='cuda'\n    )\n    topk.append(topk_batch.numpy())\n    prob.append(prob_batch.numpy())\ntopk = np.concatenate(topk)\nprob = np.concatenate(prob)\nprint(f\"topk shape: {topk.shape}, prob shape: {prob.shape}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# convert result to long-form data frame\ntest_ids = np.repeat(np.arange(n_test, dtype=np.int32), N_NEIGHBORS)\nterm_ids = topk.ravel().astype(np.int32)\nscores = prob.ravel().astype(np.float32)\n\nknn_df = pd.DataFrame({\n    \"test_id\": test_ids,\n    \"term_id\": term_ids,\n    \"score\": scores\n})\nind2term = {i:k for i,k in enumerate(vocab)}\nind2prot = {i:k for i,k in enumerate(test_df.prot_id)}\nknn_df['term'] = knn_df.term_id.map(ind2term)\nknn_df['prot_id'] = knn_df.test_id.map(ind2prot)\nknn_df = knn_df[['prot_id', 'term', 'score']]\nknn_df","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# remove negative links\nimport gc\nimport os\n\n# ==========================================\n# 1. OPTIMIZED OBO PARSER (Pre-compute Descendants)\n# ==========================================\ndef parse_obo_and_get_descendants(go_obo_path):\n    \"\"\"\n    Parses OBO and returns a dictionary mapping every GO term \n    to a set of ALL its descendants (transitive closure).\n    \"\"\"\n    print(f\"[1/5] Parsing OBO Ontology from {go_obo_path}...\")\n    children = defaultdict(set)\n    \n    if not os.path.exists(go_obo_path):\n        print(\"Warning: OBO file not found.\")\n        return {}\n\n    with open(go_obo_path, \"r\") as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == \"[Term]\":\n                cur_id = None\n            elif line.startswith(\"id: \"):\n                cur_id = line.split(\"id: \")[1].strip()\n            elif line.startswith(\"is_a: \"):\n                # Parent logic: if \"id\" is_a \"pid\", then \"id\" is a child of \"pid\"\n                pid = line.split()[1].strip()\n                if cur_id:\n                    children[pid].add(cur_id)\n            elif line.startswith(\"relationship: part_of \"):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id:\n                        children[pid].add(cur_id)\n\n    # Pre-compute transitive closure (all descendants) for every node\n    # Using a cache to avoid re-calculating common paths\n    print(f\"[1/5] Calculating transitive closure for {len(children)} parent nodes...\")\n    descendants_cache = {}\n\n    def get_all_descendants(term):\n        if term in descendants_cache:\n            return descendants_cache[term]\n        \n        direct_children = children.get(term, set())\n        all_desc = set(direct_children)\n        \n        for child in direct_children:\n            all_desc |= get_all_descendants(child)\n            \n        descendants_cache[term] = all_desc\n        return all_desc\n\n    # Force computation for all keys\n    # We iterate list(children.keys()) to avoid runtime dict change errors\n    for term in list(children.keys()):\n        get_all_descendants(term)\n        \n    print(f\"[1/5] OBO Parsed. Knowledge base ready.\")\n    return descendants_cache\n\n# ==========================================\n# 2. LOAD & MERGE NEGATIVE ANNOTATIONS\n# ==========================================\ndef get_negative_pairs(file_paths, descendants_map):\n    \"\"\"\n    Loads multiple GOA files, filters for 'NOT' qualifier, \n    propagates negatives to descendants, and returns a set of unique keys.\n    Params:\n        file_paths: list of file paths \n        descendants_map: a dictionary mapping from Go term to all of its descendants.\n    Returns: \n        A set of keys\n        \n    \"\"\"\n    unique_negatives = set()\n    \n    for path in file_paths:\n        if not os.path.exists(path):\n            continue\n            \n        print(f\"[2/5] Processing annotations: {os.path.basename(path)}\")\n        # Determine separator based on file extension\n        sep = '\\t' if path.endswith('.tsv') else ','\n        \n        # specific loading for efficiency\n        df = pd.read_csv(path, sep=sep, dtype=str, usecols=['protein_id', 'go_term', 'qualifier'])\n        \n        # Filter for Negative qualifiers (NOT)\n        neg_df = df[df['qualifier'].str.contains('NOT', na=False)]\n        print(f\"      Found {len(neg_df)} explicit negative annotations.\")\n        \n        # Group by Protein to speed up propagation\n        # protein -> [term1, term2, ...]\n        prot_to_terms = neg_df.groupby('protein_id')['go_term'].apply(list).to_dict()\n        \n        count = 0\n        for protein, terms in tqdm(prot_to_terms.items(), desc=\"Propagating Negatives\"):\n            for term in terms:\n                # 1. The term itself is negative\n                unique_negatives.add(f\"{protein}_{term}\")\n                \n                # 2. All descendants of the term are negative (Negative Propagation Rule)\n                # If protein is NOT involved in 'Metabolism', it is NOT involved in 'Gluconeogenesis'\n                if term in descendants_map:\n                    for desc in descendants_map[term]:\n                        unique_negatives.add(f\"{protein}_{desc}\")\n        \n        del df, neg_df, prot_to_terms\n        gc.collect()\n\n    return unique_negatives\n\n# ==========================================\n# 3. MAIN PIPELINE\n# ==========================================\n\n# --- Configuration ---\nOBO_PATH = \"data/Train/go-basic.obo\"\nGOA_FILES = [\n    'go_annotation/goa_uniprot_ver227.csv',\n    'go_annotation/goa_uniprot_ver228.tsv'\n]\nCHUNK_SIZE = 500_000\n\n# --- Execution ---\n\n# Step 1: Get Ontology Structure\ndescendants_map = parse_obo_and_get_descendants(OBO_PATH)\n\n# Step 2: Build Global Negative Dictionary\nnegative_keys = get_negative_pairs(GOA_FILES, descendants_map)\nprint(f\"[+] Total unique Negative Protein-GO pairs (Propagated): {len(negative_keys)}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# filtering\nknn_df['key'] = knn_df.prot_id + \"_\" + knn_df.term\nknn_df = knn_df[~knn_df.key.isin(negative_keys)]\nknn_df = knn_df.drop(columns=['key'])\nknn_df.shape","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# save to file\nknn_df.to_csv('submission.tsv', index=False, sep='\\t', header=False)\nprint(\"Submission file saved to disk!\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# making prediction using detected thresholds\ntest_probs = predict_all(model, X_test)\ny_preds = (test_probs > best_thresholds)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# get term from index\nindex2term = {int(i): t for i, t in enumerate(vocab)}\nprint(f\"No. terms: {len(index2term)}\")\nx_indx, y_indx = np.where(y_preds==1)\nrows = []\nfor i in tqdm(range(len(x_indx))):\n    row = [\n        ind2prot[x_indx[i]],\n        index2term[y_indx[i]],\n        test_probs[x_indx[i], y_indx[i]]\n    ]\n    rows.append(\n        row\n    )\nout_df = pd.DataFrame(rows, columns=['prot_id', 'term', 'score'])\nout_df","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# filtering\nout_df['key'] = out_df.prot_id + \"_\" + out_df.term\nout_df = out_df[~out_df.key.isin(negative_keys)]\nout_df = out_df.drop(columns=['key'])\nout_df.shape","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# save to file\nout_df.to_csv('submission_thresholded.tsv', index=False, sep='\\t', header=False)\nprint(\"Submission file saved to disk!\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{},"outputs":[],"execution_count":null}]}