{"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":13431864,"sourceType":"datasetVersion","datasetId":8525294}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **STAGE 0: INSTALLS AND IMPORTS** `","metadata":{}},{"cell_type":"code","source":"\n!pip install -q transformers obonet biopython accelerate datasets\n!pip install -q biopython obonet transformers\n!pip install -q pyarrow==21.0.0 --quiet\n\nimport optuna\n\nimport logging\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport os, random, time\nfrom pathlib import Path\nfrom collections import defaultdict\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom scipy import sparse\n\nimport torch\nimport torch.nn as nn\nfrom torch import optim\nfrom torch.utils.data import Dataset, DataLoader, Subset\n\nimport obonet\nfrom Bio import SeqIO\nfrom torch.cuda.amp import GradScaler, autocast\nfrom transformers import AutoTokenizer, AutoModel\n\n\nfrom sklearn.preprocessing import MultiLabelBinarizer, StandardScaler\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score, precision_score, recall_score\n\n\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"  # suppress TensorFlow/XLA logs (0=all, 3=errors only)\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\nos.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"0\"\nos.environ[\"TRANSFORMERS_NO_ADVISORY_WARNINGS\"] = \"1\"\nos.environ[\"HF_HUB_DISABLE_SYMLINKS_WARNING\"] = \"1\"\nos.environ[\"HF_HUB_REQUEST_TIMEOUT\"] = \"120\" # Environment variables for Hugging Face\nos.environ[\"HF_HUB_CONNECT_RETRIES\"] = \"10\"\nos.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n\nlogging.getLogger(\"pip\").setLevel(logging.ERROR)\nlogging.getLogger(\"transformers\").setLevel(logging.ERROR)\nlogging.getLogger(\"torch\").setLevel(logging.ERROR)\n\nprint(\"Imports OK. Torch version:\", torch.__version__)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:32:39.135161Z","iopub.execute_input":"2025-10-21T03:32:39.135752Z","iopub.status.idle":"2025-10-21T03:33:02.636522Z","shell.execute_reply.started":"2025-10-21T03:32:39.135726Z","shell.execute_reply":"2025-10-21T03:33:02.63578Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 1:  PATHS AND SETTINGS**","metadata":{}},{"cell_type":"code","source":"BASE = \"/kaggle/input/cafa-6-protein-function-prediction\"\nTRAIN_DIR = f\"{BASE}/Train\"\nTEST_DIR = f\"{BASE}/Test\"\n\nTRAIN_FASTA = f\"{TRAIN_DIR}/train_sequences.fasta\"\nTRAIN_TERMS = f\"{TRAIN_DIR}/train_terms.tsv\"\nTRAIN_TAX = f\"{TRAIN_DIR}/train_taxonomy.tsv\"\nGO_OBO = f\"{TRAIN_DIR}/go-basic.obo\"\n\nTEST_FASTA = f\"{TEST_DIR}/testsuperset.fasta\"\nTEST_TAXON = f\"{TEST_DIR}/testsuperset-taxon-list.tsv\"\n\nIA_PATH = f\"{BASE}/IA.tsv\"\nSAMPLE_SUB = f\"{BASE}/sample_submission.tsv\"\n\nEMBED_CACHE_TRAIN = \"/kaggle/input/embeddings-cache/embeddings_cache/train_esm2_embeddings.parquet\"\nEMBED_CACHE_TEST = \"/kaggle/input/embeddings-cache/embeddings_cache/test_esm2_embeddings.parquet\"\n\nOUT_DIR = \"/kaggle/working/cafa_mlp_output\"\nos.makedirs(OUT_DIR, exist_ok=True)\n\n# Experiment switches (tune for debugging / full run)\nSEED = 42\nEMBED_MODEL = \"facebook/esm2_t6_8M_UR50D\"   # balanced speed & quality\nBATCH_SIZE_EMBED = 4    # embedding batch size (reduce if OOM)\nBATCH_SIZE_TRAIN = 128   # training batch size (reduce if OOM)\nMAX_SEQS_TO_EMBED = None # None => embed all train sequences. Set small for debugging.\nMAX_LABELS = None        # None => use all GO terms. Use int (e.g., 5000) to limit labels while debugging.\n\nHPO_TRAIN_FRAC = 0.15  # use 15% of training data for speed\nHPO_VAL_FRAC = 0.25    # use 25% of validation data\n\nEPOCHS = 20\nLR = 1e-4\nWEIGHT_DECAY = 1e-5\nPATIENCE = 3\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", DEVICE)\nprint(\"Paths set. Output dir:\", OUT_DIR)\n\nprint(\"🔧 Experiment Settings\")\nprint(\"-\"*40)\nprint(f\"Random seed: {SEED}\")\nprint(f\"ESM2 model: {EMBED_MODEL}\")\nprint(f\"Embedding batch size: {BATCH_SIZE_EMBED}\")\nprint(f\"Training batch size: {BATCH_SIZE_TRAIN}\")\nprint(f\"Max train sequences to embed: {MAX_SEQS_TO_EMBED}\")\nprint(f\"Max GO terms (labels): {MAX_LABELS}\")\nprint(f\"Training epochs: {EPOCHS}\")\nprint(f\"Learning rate: {LR}\")\nprint(f\"Weight decay: {WEIGHT_DECAY}\")\nprint(f\"Early stopping patience: {PATIENCE}\")\nprint(f\"Training Fraction: {HPO_TRAIN_FRAC}\")\nprint(f\"Validation Fraction: {HPO_VAL_FRAC}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:02.638092Z","iopub.execute_input":"2025-10-21T03:33:02.638573Z","iopub.status.idle":"2025-10-21T03:33:02.729104Z","shell.execute_reply.started":"2025-10-21T03:33:02.638553Z","shell.execute_reply":"2025-10-21T03:33:02.727993Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 2: LOAD METADATA AND SEQUENCES**","metadata":{}},{"cell_type":"code","source":"# Load dataset files (terms, taxonomy, IA, sequences)\ndef seed_everything(seed=SEED):\n    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)\n    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)\n\nseed_everything()\n\n# Load metadata \ntrain_terms = pd.read_csv(TRAIN_TERMS, sep=\"\\t\", header=None,\n                          names=[\"protein_id\",\"go_term\",\"ontology\"], dtype=str)\ntrain_terms = train_terms[~train_terms['protein_id'].str.contains('Entry', na=False)].reset_index(drop=True)\n\ntrain_tax = pd.read_csv(TRAIN_TAX, sep=\"\\t\", header=None,\n                        names=[\"protein_id\",\"taxon_id\"], dtype=str)\n\nia = pd.read_csv(IA_PATH, sep=\"\\t\", header=None,\n                 names=[\"go_term\",\"ia_weight\"], dtype={'go_term':str,'ia_weight':float})\n \ngo_graph = obonet.read_obo(GO_OBO)\n\n# Load FASTA sequences \ndef load_fasta_dict(path, max_records=None):\n    out = {}\n    for i, rec in enumerate(SeqIO.parse(path, \"fasta\")):\n        prot_id = rec.id.split(\"|\")[1] if \"|\" in rec.id else rec.id\n        out[prot_id] = str(rec.seq)\n        if max_records and (i+1) >= max_records:\n            break\n    return out\n\ntrain_seqs = load_fasta_dict(TRAIN_FASTA, max_records=MAX_SEQS_TO_EMBED)\ntest_seqs = load_fasta_dict(TEST_FASTA, max_records=None)\n\n# Merge annotations with taxonomy \ntrain_df = train_terms.merge(train_tax, on=\"protein_id\", how=\"left\")\n\n# -------------------- Dataset summary --------------------\nprint(\"\\n📊 Dataset Overview\\n\" + \"-\"*40)\nprint(f\"Total GO ontology terms: {len(go_graph)}\")\nprint(f\"Number of training sequences: {len(train_seqs)}\")\nprint(f\"Number of test sequences: {len(test_seqs)}\")\nprint(f\"Training DataFrame shape: {train_df.shape}\\n\")\n\nprint(\"Summary of key datasets:\")\nprint(f\"• train_terms dataframe: {train_terms.shape}\")\nprint(f\"• train_taxonomy dataframe: {train_tax.shape}\")\nprint(f\"• Information content (IA) table: {ia.shape}\")\nprint(f\"• Number of training sequences: {len(train_seqs)}\")\nprint(f\"• Number of test sequences: {len(test_seqs)}\")\nprint(f\"• Total GO terms in ontology graph: {len(go_graph)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:02.730186Z","iopub.execute_input":"2025-10-21T03:33:02.730517Z","iopub.status.idle":"2025-10-21T03:33:11.395258Z","shell.execute_reply.started":"2025-10-21T03:33:02.73049Z","shell.execute_reply":"2025-10-21T03:33:11.394381Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 3:  EMBEDDING EXTRACTION (ESM2) WITH CACHING**","metadata":{}},{"cell_type":"code","source":"def load_fasta_sequences(file_path):\n    sequences = {}\n    for record in SeqIO.parse(file_path, \"fasta\"):\n        prot_id = record.id.split(\"|\")[1] if \"|\" in record.id else record.id\n        sequences[prot_id] = str(record.seq)\n    return sequences\n\ntrain_sequences = load_fasta_sequences(TRAIN_FASTA)\ntest_sequences = load_fasta_sequences(TEST_FASTA)\n\ndef extract_embeddings(seqs_dict, cache_path, model_name=EMBED_MODEL, batch_size=BATCH_SIZE_EMBED, device=DEVICE):\n    cache = Path(cache_path)\n    if cache.exists():\n        print(\"Loading cached embeddings:\", cache_path)\n        return pd.read_parquet(cache_path)\n\n    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)\n    model = AutoModel.from_pretrained(model_name).to(device)\n    model.eval()\n\n    prot_ids = list(seqs_dict.keys())\n    embeddings = []\n    ids = []\n\n    for i in tqdm(range(0, len(prot_ids), batch_size), desc=\"Embedding batches\"):\n        batch_ids = prot_ids[i:i+batch_size]\n        batch_seqs = [seqs_dict[x] for x in batch_ids]\n        inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=1022)\n        inputs = {k: v.to(device) for k, v in inputs.items()}\n\n        with torch.no_grad():\n            out = model(**inputs)\n            emb = out.last_hidden_state.mean(dim=1).cpu().numpy()\n\n        embeddings.append(emb)\n        ids.extend(batch_ids)\n        torch.cuda.empty_cache()\n\n    emb_mat = np.vstack(embeddings)\n    df = pd.DataFrame(emb_mat)\n    df.insert(0, \"protein_id\", ids)\n    df.to_parquet(cache_path, index=False)\n    print(\"Saved embeddings to:\", cache_path, \"shape:\", df.shape)\n    return df\n\ntrain_emb_df = extract_embeddings(train_sequences, EMBED_CACHE_TRAIN)\ntest_emb_df = extract_embeddings(test_sequences, EMBED_CACHE_TEST)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:11.397012Z","iopub.execute_input":"2025-10-21T03:33:11.397305Z","iopub.status.idle":"2025-10-21T03:33:16.233096Z","shell.execute_reply.started":"2025-10-21T03:33:11.397284Z","shell.execute_reply":"2025-10-21T03:33:16.232394Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 4: PREPARE MULTI-LABEL TARGETS**","metadata":{}},{"cell_type":"code","source":"# Build label vocabulary and Y matrix\n# Optionally limit labels with MAX_LABELS during debugging\nall_go_terms = sorted(train_terms['go_term'].unique().tolist())\nif MAX_LABELS:\n    print(\"Limiting labels to:\", MAX_LABELS)\n    all_go_terms = all_go_terms[:MAX_LABELS]\nprint(\"Using #labels:\", len(all_go_terms))\n\n# IA weights array (for later)\nia_map = dict(zip(ia['go_term'], ia['ia_weight']))\nia_weights = np.array([ia_map.get(g, 0.0) for g in all_go_terms], dtype=float)\n\n# Build grouped labels per protein\ntt = train_terms[train_terms['go_term'].isin(all_go_terms)]\nlabels_grouped = tt.groupby('protein_id')['go_term'].apply(list).reset_index()\n\n# Keep only proteins with embeddings\navailable_train_ids = set(train_emb_df['protein_id'].tolist())\nlabels_grouped = labels_grouped[labels_grouped['protein_id'].isin(available_train_ids)].reset_index(drop=True)\nprint(\"Proteins with labels & embeddings:\", labels_grouped.shape[0])\n\n# Make label matrix\nmlb = MultiLabelBinarizer(classes=all_go_terms)\nY = mlb.fit_transform(labels_grouped['go_term'])\nprint(\"Label matrix shape:\", Y.shape)\n\n# Build X aligned to labels\nemb_map = train_emb_df.set_index('protein_id')\nX_rows = []\nfor pid in labels_grouped['protein_id']:\n    row = emb_map.loc[pid].drop(labels=['protein_id'], errors='ignore') if 'protein_id' in emb_map.columns else emb_map.loc[pid]\n    X_rows.append(row.values)\nX = np.vstack(X_rows)\nprint(\"Final X shape:\", X.shape)\n\n# -----------------------------------------\n# ✅ Normalize embeddings\n# Step 1: L2 normalization per protein vector\nX = X / np.linalg.norm(X, axis=1, keepdims=True)\n\n# Step 2: StandardScaler normalization across features\nscaler = StandardScaler()\nX = scaler.fit_transform(X)\n\nprint(\"✅ Embeddings normalized.\")\nprint(\"Mean (overall):\", np.mean(X))\nprint(\"Std (overall):\", np.std(X))\nprint(\"Final X shape after normalization:\", X.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:16.233898Z","iopub.execute_input":"2025-10-21T03:33:16.234149Z","iopub.status.idle":"2025-10-21T03:33:21.320774Z","shell.execute_reply.started":"2025-10-21T03:33:16.234125Z","shell.execute_reply":"2025-10-21T03:33:21.319891Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 5: DATASET AND DATALOADER**","metadata":{}},{"cell_type":"code","source":"class ProteinDataset(Dataset):\n    def __init__(self, X_np, Y_np):\n        self.X = torch.from_numpy(X_np).float()\n        self.Y = torch.from_numpy(Y_np).float()\n    def __len__(self): return len(self.X)\n    def __getitem__(self, idx): return self.X[idx], self.Y[idx]\n\n# Split train/val\ntrain_idx, val_idx = train_test_split(np.arange(len(X)), test_size=0.15, random_state=SEED, shuffle=True)\ntrain_ds = ProteinDataset(X[train_idx], Y[train_idx])\nval_ds = ProteinDataset(X[val_idx], Y[val_idx])\ntrain_loader = DataLoader(train_ds, batch_size=BATCH_SIZE_TRAIN, shuffle=True, num_workers=2, pin_memory=True)\nval_loader = DataLoader(val_ds, batch_size=BATCH_SIZE_TRAIN, shuffle=False, num_workers=2, pin_memory=True)\n\nprint(\"Train/Val sizes:\", len(train_ds), len(val_ds))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:21.321819Z","iopub.execute_input":"2025-10-21T03:33:21.322157Z","iopub.status.idle":"2025-10-21T03:33:35.438443Z","shell.execute_reply.started":"2025-10-21T03:33:21.322127Z","shell.execute_reply":"2025-10-21T03:33:35.437653Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 5: MODEL, LOSS, EVALUATE, HPO**","metadata":{}},{"cell_type":"code","source":"# Deep MLP Model\nclass DeepMLP(nn.Module):\n    def __init__(self, input_dim, num_labels, hidden_dims=[1024, 512, 256], dropout=0.3):\n        super().__init__()\n        layers = []\n        prev = input_dim\n        for h in hidden_dims:\n            layers += [\n                nn.Linear(prev, h),\n                nn.BatchNorm1d(h),\n                nn.ReLU(inplace=True),\n                nn.Dropout(dropout)\n            ]\n            prev = h\n        layers.append(nn.Linear(prev, num_labels))\n        self.net = nn.Sequential(*layers)\n    def forward(self, x):\n        return self.net(x)\n\n\n# Focal Loss\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=1.0, gamma=2.0, reduction='mean'):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.reduction = reduction\n        self.bce = nn.BCEWithLogitsLoss(reduction='none')\n\n    def forward(self, logits, targets):\n        bce_loss = self.bce(logits, targets)\n        probs = torch.sigmoid(logits)\n        focal_weight = self.alpha * (\n            (1 - probs) ** self.gamma * targets + (probs ** self.gamma) * (1 - targets)\n        )\n        loss = focal_weight * bce_loss\n        return loss.mean() if self.reduction == 'mean' else loss.sum()\n\n\n# Faster Evaluation Function (vectorized)\ndef evaluate(loader, model, threshold=0.25):\n    model.eval()\n    y_true_list, y_pred_list = [], []\n    with torch.no_grad():\n        for xb, yb in loader:\n            xb = xb.to(DEVICE)\n            probs = torch.sigmoid(model(xb)).cpu()\n            preds = (probs >= threshold).float()\n            y_pred_list.append(preds)\n            y_true_list.append(yb)\n    y_true = torch.cat(y_true_list).numpy()\n    y_pred = torch.cat(y_pred_list).numpy()\n    return (\n        f1_score(y_true, y_pred, average='micro', zero_division=0),\n        precision_score(y_true, y_pred, average='micro', zero_division=0),\n        recall_score(y_true, y_pred, average='micro', zero_division=0),\n    )\n\n\n# HPO Optimization Function\ndef objective(trial):\n    # Suggested hyperparameters\n    hidden_dims = [\n        trial.suggest_categorical(\"h1\", [512, 1024, 2048]),\n        trial.suggest_categorical(\"h2\", [256, 512, 1024])\n    ]\n    dropout = trial.suggest_float(\"dropout\", 0.1, 0.5)\n    lr = trial.suggest_float(\"lr\", 1e-5, 1e-3, log=True)\n\n    # Build model\n    model = DeepMLP(\n        input_dim=X.shape[1],\n        num_labels=Y.shape[1],\n        hidden_dims=hidden_dims,\n        dropout=dropout\n    ).to(DEVICE)\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=WEIGHT_DECAY)\n    criterion = FocalLoss(alpha=1.0, gamma=2.0)\n\n    # Train for fewer epochs for speed\n    best_val_f1 = 0.0\n    for epoch in range(2):  # Reduced from 3 to 2\n        model.train()\n        for xb, yb in train_loader:\n            xb, yb = xb.to(DEVICE), yb.to(DEVICE)\n            optimizer.zero_grad(set_to_none=True)\n            loss = criterion(model(xb), yb)\n            loss.backward()\n            optimizer.step()\n\n        # Validate\n        val_f1, _, _ = evaluate(val_loader, model)\n        trial.report(val_f1, step=epoch)\n\n        # Optuna pruning check\n        if trial.should_prune():\n            raise optuna.exceptions.TrialPruned()\n\n        best_val_f1 = max(best_val_f1, val_f1)\n\n    return best_val_f1\n\n\n# Optuna study setup\n#study = optuna.create_study( study_name=\"protein_mlp_hpo_fast\",  direction=\"maximize\",   pruner=optuna.pruners.MedianPruner(n_startup_trials=2, n_warmup_steps=1, interval_steps=1))\n# set n_trials 10/20 before execution\n\n# Reduce number of trials for faster optimization\n#study.optimize(objective, n_trials=0, show_progress_bar=True)\n\n#print(\"Best hyperparameters found:\")\n#print(study.best_params)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:35.439279Z","iopub.execute_input":"2025-10-21T03:33:35.439598Z","iopub.status.idle":"2025-10-21T03:33:43.189372Z","shell.execute_reply.started":"2025-10-21T03:33:35.439563Z","shell.execute_reply":"2025-10-21T03:33:43.188229Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 7: TRAINING (EARLY STOPPING)**","metadata":{}},{"cell_type":"code","source":"# [I 2025-10-20 17:24:12,504] Trial 1 finished with value: 0.16273737993294707 and parameters: {'h1': 2048, 'h2': 256, 'dropout': 0.1889064287439644, 'lr': 0.00011241516730960752}. Best is trial 1 with value: 0.16273737993294707.\n# Use best hyperparameters from study\n# best_params = study.best_params\nhidden_dims = [2048, 256]\ndropout = 0.1889064287439644\nlr = 0.00011241516730960752\n\n# Build final model\nmodel = DeepMLP(input_dim=X.shape[1], num_labels=Y.shape[1], hidden_dims=hidden_dims, dropout=dropout).to(DEVICE)\ncriterion = FocalLoss(alpha=1.0, gamma=2.0)\noptimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=WEIGHT_DECAY)\n\n# Train full model\nbest_f1, patience = 0, 0\nfor epoch in range(1, EPOCHS + 1):\n    model.train()\n    train_loss = 0.0\n    for xb, yb in tqdm(train_loader, desc=f\"Epoch {epoch}/{EPOCHS}\"):\n        xb, yb = xb.to(DEVICE), yb.to(DEVICE)\n        optimizer.zero_grad()\n        with autocast():\n            loss = criterion(model(xb), yb)\n        loss.backward()\n        optimizer.step()\n        train_loss += loss.item()\n    train_loss /= len(train_loader)\n\n    val_f1, val_prec, val_rec = evaluate(val_loader, model)\n    print(f\"Epoch {epoch}: loss={train_loss:.6f} val_f1={val_f1:.4f} prec={val_prec:.4f} rec={val_rec:.4f}\")\n\n    if val_f1 > best_f1:\n        best_f1 = val_f1\n        patience = 0\n        torch.save(model.state_dict(), os.path.join(OUT_DIR, \"best_mlp.pt\"))\n    else:\n        patience += 1\n    if patience >= PATIENCE:\n        print(\"Early stopping triggered.\")\n        break\n\nprint(\"✅ Best val F1:\", best_f1)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:51.589213Z","iopub.execute_input":"2025-10-21T03:33:51.589869Z","iopub.status.idle":"2025-10-21T03:39:17.61352Z","shell.execute_reply.started":"2025-10-21T03:33:51.58983Z","shell.execute_reply":"2025-10-21T03:39:17.612282Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# **STAGE 8: TESTING AND PREPARE SUBMMISSION**","metadata":{}},{"cell_type":"code","source":"# Build test matrix aligned with training embedding columns\ntest_ids = test_emb_df['protein_id'].tolist()\nX_test = test_emb_df.drop(columns=['protein_id']).values\ntest_ds = ProteinDataset(X_test, np.zeros((X_test.shape[0], num_labels)))\ntest_loader = DataLoader(test_ds, batch_size=BATCH_SIZE_TRAIN, shuffle=False, num_workers=2)\n\nall_probs = []\nmodel.eval()\nwith torch.no_grad():\n    for xb, _ in tqdm(test_loader, desc=\"Predict test\"):\n        xb = xb.to(DEVICE)\n        logits = model(xb)\n        probs = torch.sigmoid(logits).cpu().numpy()\n        all_probs.append(probs)\nall_probs = np.vstack(all_probs)\nprint(\"Test probs shape:\", all_probs.shape)\n\n# Convert to CAFA TSV: (protein_id, go_term, score)\nTHRESH = 0.1   # adjust threshold\nTOP_K = 1500   # competition cap per protein overall\nrows = []\ngo_terms = list(mlb.classes_)\nfor i, pid in enumerate(test_ids):\n    probs = all_probs[i]\n    idxs = np.where(probs > THRESH)[0]\n    if len(idxs) == 0:\n        # fallback: top 1\n        idxs = [int(np.argmax(probs))]\n    # sort by prob desc and cap\n    idxs = idxs[np.argsort(probs[idxs])[::-1]]\n    idxs = idxs[:TOP_K]\n    for j in idxs:\n        rows.append((pid, go_terms[j], float(probs[j])))\n\nsub_df = pd.DataFrame(rows, columns=[\"protein_id\",\"go_term\",\"score\"])\nout_path = os.path.join(OUT_DIR, \"submission_mlp_raw.tsv\")\nsub_df.to_csv(out_path, sep=\"\\t\", header=False, index=False)\nprint(\"Raw submission saved:\", out_path, \"rows:\", len(sub_df))\n\n# Optional: propagate to ancestors using GO graph (recommended)\nprint(\"Loading GO graph and propagating predictions ...\")\ngo_graph = obonet.read_obo(GO_OBO)\nparents = {n: list(go_graph.predecessors(n)) for n in go_graph.nodes()}\n\ndef propagate(term_scores, parents_map):\n    propagated = dict(term_scores)\n    for term, score in list(term_scores.items()):\n        stack = [term]\n        while stack:\n            t = stack.pop()\n            for p in parents_map.get(t, []):\n                if propagated.get(p, 0.0) < score:\n                    propagated[p] = score\n                    stack.append(p)\n    return propagated\n\nfrom collections import defaultdict\npred_by_prot = defaultdict(dict)\nfor pid, term, score in rows:\n    pred_by_prot[pid][term] = max(pred_by_prot[pid].get(term, 0.0), score)\n\nprop_rows = []\nfor pid, ts in tqdm(pred_by_prot.items(), desc=\"Propagate\"):\n    pr = propagate(ts, parents)\n    sorted_terms = sorted(pr.items(), key=lambda x: x[1], reverse=True)[:TOP_K]\n    for t,s in sorted_terms:\n        prop_rows.append((pid,t,s))\n\nprop_df = pd.DataFrame(prop_rows, columns=[\"protein_id\",\"go_term\",\"score\"])\nprop_out = os.path.join(OUT_DIR, \"submission_mlp_propagated.tsv\")\nprop_df.to_csv(prop_out, sep=\"\\t\", header=False, index=False)\nprint(\"Propagated submission saved:\", prop_out, \"rows:\", len(prop_df))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-21T03:33:43.192245Z","iopub.status.idle":"2025-10-21T03:33:43.192516Z","shell.execute_reply.started":"2025-10-21T03:33:43.192402Z","shell.execute_reply":"2025-10-21T03:33:43.192415Z"}},"outputs":[],"execution_count":null}]}