{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14089785,"sourceType":"datasetVersion","datasetId":8971485},{"sourceId":276367826,"sourceType":"kernelVersion"},{"sourceId":285164560,"sourceType":"kernelVersion"}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --no-index --find-links=/kaggle/input/cafa6-bio-packages/bio_packages Bio==1.8.1\n!pip install --no-index --find-links=/kaggle/input/cafa6-bio-packages/bio_packages protobuf==3.20.3\n!pip install --no-index --find-links=/kaggle/input/cafa6-bio-packages/bio_packages obonet==1.1.1","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:57:55.528218Z","iopub.execute_input":"2025-12-10T11:57:55.528613Z","iopub.status.idle":"2025-12-10T11:58:07.699593Z","shell.execute_reply.started":"2025-12-10T11:57:55.528576Z","shell.execute_reply":"2025-12-10T11:58:07.698861Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, gc, math, json\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm\n\nfrom Bio import SeqIO\nimport obonet  # <-- new\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score, precision_recall_curve\n\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\nfrom einops import reduce\nfrom torch import nn\nfrom torch.utils.data import TensorDataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:07.701151Z","iopub.execute_input":"2025-12-10T11:58:07.701429Z","iopub.status.idle":"2025-12-10T11:58:18.169791Z","shell.execute_reply.started":"2025-12-10T11:58:07.701376Z","shell.execute_reply":"2025-12-10T11:58:18.169206Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"DATA_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\n\nFP_TRAIN_FASTA = f\"{DATA_DIR}/Train/train_sequences.fasta\"\nFP_TEST_FASTA  = f\"{DATA_DIR}/Test/testsuperset.fasta\"\nFP_TERMS       = f\"{DATA_DIR}/Train/train_terms.tsv\"\nFP_TAX         = f\"{DATA_DIR}/Train/train_taxonomy.tsv\"\nFP_IA          = f\"{DATA_DIR}/IA.tsv\"            \nFP_OBO         = f\"{DATA_DIR}/Train/go-basic.obo\"      \n\nMAX_LEN = 1024   ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:18.170452Z","iopub.execute_input":"2025-12-10T11:58:18.170902Z","iopub.status.idle":"2025-12-10T11:58:18.175503Z","shell.execute_reply.started":"2025-12-10T11:58:18.170876Z","shell.execute_reply":"2025-12-10T11:58:18.174806Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_fasta(path):\n    \"\"\"Parse FASTA to dict: protein_id -> sequence.\"\"\"\n    seqs = {}\n    with open(path) as handle:\n        for rec in SeqIO.parse(handle, \"fasta\"):\n            pid = rec.id.split('|')[1] if '|' in rec.id else rec.id\n            seqs[pid] = str(rec.seq)\n    return seqs\n\ntrain_seqs = load_fasta(FP_TRAIN_FASTA)\ntest_seqs  = load_fasta(FP_TEST_FASTA)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:18.176241Z","iopub.execute_input":"2025-12-10T11:58:18.177221Z","iopub.status.idle":"2025-12-10T11:58:20.218773Z","shell.execute_reply.started":"2025-12-10T11:58:18.177202Z","shell.execute_reply":"2025-12-10T11:58:20.217992Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms = pd.read_csv(FP_TERMS, sep=\"\\t\", names=[\"protein_id\",\"go_id\",\"ontology\"], skiprows=1)\ntrain_tax   = pd.read_csv(FP_TAX,   sep=\"\\t\", names=[\"protein_id\",\"taxon_id\"])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:20.220549Z","iopub.execute_input":"2025-12-10T11:58:20.220769Z","iopub.status.idle":"2025-12-10T11:58:20.517270Z","shell.execute_reply.started":"2025-12-10T11:58:20.220751Z","shell.execute_reply":"2025-12-10T11:58:20.516720Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------\n# GO graph based label filtering\n# -------------------------\n\ndef find_best_filter_from_graph(go_graph, train_terms_df, target_labels=1500,\n                                max_children_grid=None, min_count_grid=None):\n    \"\"\"\n    Use GO graph + term frequency to find (max_children, min_count)\n    which keeps a number of labels close to target_labels.\n    \"\"\"\n    # Reverse graph: parent -> children\n    go_rev = go_graph.reverse(copy=True)\n    child_counts = {n: go_rev.out_degree(n) for n in go_rev.nodes()}\n\n    # Default grids\n    if max_children_grid is None:\n        max_deg = max(child_counts.values()) if child_counts else 1\n        max_children_grid = np.unique(\n            np.round(np.logspace(0, np.log10(max_deg), 15)).astype(int)\n        )\n    if min_count_grid is None:\n        min_count_grid = np.arange(0, 201, 10)\n\n    # Build GO term DataFrame: term + number of children\n    go_df = pd.DataFrame(\n        [(n, child_counts[n]) for n in go_graph.nodes()],\n        columns=[\"go_id\", \"n_children\"]\n    )\n\n    # Frequency of each term in training annotations\n    term_counts = train_terms_df[\"go_id\"].value_counts()\n\n    results = []\n    for max_children in max_children_grid:\n        semi_leaf_terms = {n for n, c in child_counts.items() if c <= max_children}\n        for min_count in min_count_grid:\n            frequent_terms = set(term_counts[term_counts >= min_count].index)\n            final_terms = semi_leaf_terms & frequent_terms\n            n_labels = len(final_terms)\n            results.append((max_children, min_count, n_labels))\n\n    summary_df = pd.DataFrame(results, columns=[\"max_children\", \"min_count\", \"n_labels\"])\n    summary_df[\"diff\"] = (summary_df[\"n_labels\"] - target_labels).abs()\n    best_row = summary_df.loc[summary_df[\"diff\"].idxmin()]\n\n    best_max_children = int(best_row[\"max_children\"])\n    best_min_count = int(best_row[\"min_count\"])\n    best_n_labels = int(best_row[\"n_labels\"])\n\n    print(f\"Target labels: {target_labels}\")\n    print(f\"Best params: max_children={best_max_children}, min_count={best_min_count} -> {best_n_labels} labels\")\n\n    return best_max_children, best_min_count, best_n_labels, summary_df, go_df\n\n\ndef filter_train_terms(train_terms_df, go_df, max_children, min_count):\n    \"\"\"\n    Keep only GO terms which are:\n      - close to leaves in GO graph (n_children <= max_children)\n      - frequent enough in training annotations (count >= min_count)\n    \"\"\"\n    term_counts = train_terms_df[\"go_id\"].value_counts()\n\n    semi_leaf_terms = set(go_df.loc[go_df[\"n_children\"] <= max_children, \"go_id\"])\n    frequent_terms = set(term_counts[term_counts >= min_count].index)\n\n    selected_terms = semi_leaf_terms & frequent_terms\n\n    filtered_train_terms_df = train_terms_df[train_terms_df[\"go_id\"].isin(selected_terms)].copy()\n\n    print(f\"Original annotations: {len(train_terms_df)}\")\n    print(f\"Filtered annotations: {len(filtered_train_terms_df)}\")\n    print(f\"Unique GO terms before filtering: {train_terms_df['go_id'].nunique()}\")\n    print(f\"Unique GO terms after filtering:  {filtered_train_terms_df['go_id'].nunique()}\")\n\n    return filtered_train_terms_df, selected_terms\n\n\n# ---- Load GO graph and apply filtering ----\ngo_graph = obonet.read_obo(FP_OBO)\n\n# You can tune target_labels: e.g. 500 / 1000 / 2000 depending on resources\nbest_max_children, best_min_count, best_n_labels, summary_df, go_df = find_best_filter_from_graph(\n    go_graph, train_terms, target_labels=1500\n)\n\nfiltered_train_terms, selected_terms = filter_train_terms(\n    train_terms, go_df,\n    max_children=best_max_children,\n    min_count=best_min_count\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:20.518030Z","iopub.execute_input":"2025-12-10T11:58:20.518271Z","iopub.status.idle":"2025-12-10T11:58:27.544591Z","shell.execute_reply.started":"2025-12-10T11:58:20.518254Z","shell.execute_reply":"2025-12-10T11:58:27.543815Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df = pd.DataFrame({\n    \"protein_id\": list(train_seqs.keys()),\n    \"sequence\":   list(train_seqs.values())\n})\n\n# Only keep proteins that still have at least one retained GO term\ntrain_df = train_df.merge(filtered_train_terms, on=\"protein_id\", how=\"inner\")\ntrain_df[\"seq_len\"] = train_df[\"sequence\"].str.len()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:27.545332Z","iopub.execute_input":"2025-12-10T11:58:27.546069Z","iopub.status.idle":"2025-12-10T11:58:27.744009Z","shell.execute_reply.started":"2025-12-10T11:58:27.546047Z","shell.execute_reply":"2025-12-10T11:58:27.743255Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Train proteins:\", len(train_seqs))\nprint(\"Test proteins :\", len(test_seqs))\nprint(\"Unique GO terms:\", train_terms[\"go_id\"].nunique())\nprint(\"Ontology counts:\\n\", train_terms[\"ontology\"].value_counts(dropna=True))\nprint('---')\nprint(\"Train proteins (with at least 1 retained GO):\", train_df[\"protein_id\"].nunique())\nprint(\"Unique GO terms (after filtering):\", filtered_train_terms[\"go_id\"].nunique())\nprint(\"Ontology counts:\\n\", filtered_train_terms[\"ontology\"].value_counts(dropna=True))\n\nplt.hist(train_df[\"seq_len\"], bins=60)\nplt.title(\"Sequence Length Distribution\")\nplt.xlabel(\"Length\")\nplt.ylabel(\"Count\")\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:27.744854Z","iopub.execute_input":"2025-12-10T11:58:27.745345Z","iopub.status.idle":"2025-12-10T11:58:28.122775Z","shell.execute_reply.started":"2025-12-10T11:58:27.745302Z","shell.execute_reply":"2025-12-10T11:58:28.122099Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"protein2labels = (train_df.groupby(\"protein_id\")[\"go_id\"]\n                  .apply(lambda x: sorted(set(x.dropna()))))\n\nmlb = MultiLabelBinarizer()\nY = mlb.fit_transform(protein2labels)  # shape: [N_proteins, N_terms]\nprotein_ids = protein2labels.index.tolist()\n\nia_df = pd.read_csv(FP_IA, sep=\"\\t\", header=None, names=[\"go_id\",\"weight\"])\nia_map = dict(zip(ia_df[\"go_id\"], ia_df[\"weight\"]))\nterm_weights = np.array([ia_map.get(go, 1.0) for go in mlb.classes_], dtype=np.float32)\n\nprint(\"Targets shape:\", Y.shape)\nprint(\"IA weights coverage:\", (term_weights>0).mean())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:28.123540Z","iopub.execute_input":"2025-12-10T11:58:28.123837Z","iopub.status.idle":"2025-12-10T11:58:33.571280Z","shell.execute_reply.started":"2025-12-10T11:58:28.123820Z","shell.execute_reply":"2025-12-10T11:58:33.570478Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ntorch.manual_seed(42)\nif device.type == \"cuda\":\n    torch.cuda.manual_seed_all(42)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:33.572157Z","iopub.execute_input":"2025-12-10T11:58:33.572429Z","iopub.status.idle":"2025-12-10T11:58:33.635205Z","shell.execute_reply.started":"2025-12-10T11:58:33.572391Z","shell.execute_reply":"2025-12-10T11:58:33.634459Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model_name = \"/kaggle/input/cafa6-esm2-model/esm2_t6_8m_UR50D\"   # or \"Rostlab/prot_bert\" / \"Rostlab/prot_t5_xl_uniref50\"\ntokenizer = AutoTokenizer.from_pretrained(model_name)\nmodel  = AutoModel.from_pretrained(model_name).to(device)\nmodel.eval()\n\nH = model.config.hidden_size\n\n\n@torch.no_grad()\ndef embed_sequence(seq: str) -> np.ndarray:\n    \"\"\"Return pooled embedding vector (mean pooling over tokens).\"\"\"\n    batch = tokenizer(seq, return_tensors=\"pt\", truncation=True, max_length=MAX_LEN, add_special_tokens=True)\n    batch = {k: v.to(device) for k, v in batch.items()}\n    out = model(**batch, output_hidden_states=True)\n    last = out.hidden_states[-1]               # [B, L, H]\n    pooled = last.mean(dim=1)[0].detach().cpu().numpy()  # [H]\n    return pooled","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:57.615144Z","iopub.execute_input":"2025-12-10T11:58:57.615884Z","iopub.status.idle":"2025-12-10T11:58:57.620817Z","shell.execute_reply.started":"2025-12-10T11:58:57.615863Z","shell.execute_reply":"2025-12-10T11:58:57.620124Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"i_have_data = False  # set True if you calculated the embedding of the training set\n\nif not i_have_data:\n    X = np.zeros((len(protein_ids), H), dtype=np.float32)\n    \n    for i, pid in enumerate(tqdm(protein_ids, desc=\"Embedding trainset\")):\n        X[i] = embed_sequence(train_seqs[pid])\n    \n    np.save('esm_trainset.npy', X)  # create dataset for next time loading\nelse:\n    X = np.load('/kaggle/input/cafa6-esm-embedding-trainset/esm_trainset.npy')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:05:29.798310Z","iopub.execute_input":"2025-12-10T12:05:29.799129Z","iopub.status.idle":"2025-12-10T12:05:35.128536Z","shell.execute_reply.started":"2025-12-10T12:05:29.799105Z","shell.execute_reply":"2025-12-10T12:05:35.127454Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MLP(nn.Module):\n    def __init__(self, in_dim, out_dim):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(in_dim, 1024),\n            nn.ReLU(inplace=True),\n            nn.Dropout(0.25),\n            nn.Linear(1024, 512),\n            nn.ReLU(inplace=True),\n            nn.Dropout(0.25),\n            nn.Linear(512, out_dim)  # logits\n        )\n    def forward(self, x):\n        return self.net(x)\n\n\ndef evaluate(model, loader):\n    model.eval()\n    total_loss = 0.0\n    n = 0\n    with torch.no_grad():\n        for xb, yb in loader:\n            xb = xb.to(device, non_blocking=True)\n            yb = yb.to(device, non_blocking=True)\n            logits = model(xb)\n            loss = criterion(logits, yb) if y_tr.ndim == 2 else criterion(logits, yb.long())\n            bs = xb.size(0)\n            total_loss += loss.item() * bs\n            n += bs\n    return total_loss / max(n, 1)\n\n\ndef cafa_weighted_fmax(\n    y_true: np.ndarray,\n    y_pred: np.ndarray,\n    ia: np.ndarray,\n    thresholds: np.ndarray | None = None,\n    eps: float = 1e-12\n):\n    # Basic checks\n    assert y_true.shape == y_pred.shape, \"y_true and y_pred must have same shape\"\n    n_proteins, n_terms = y_true.shape\n    assert ia.shape[0] == n_terms, \"ia length must equal number of terms\"\n\n    if thresholds is None:\n        # Skip 0 and 1 to avoid trivial all-0/all-1 predictions\n        thresholds = np.linspace(0.0, 1.0, 101)[1:-1]\n\n    # Broadcast IA to (1, n_terms)\n    ia_row = ia.reshape(1, -1)\n\n    precisions = []\n    recalls = []\n    f1s = []\n    taus = []\n\n    for tau in tqdm(thresholds):\n        # Binary predictions at threshold tau\n        pred_bin = (y_pred >= tau).astype(np.int8)\n\n        # Weighted confusion matrix components\n        tp = ((pred_bin & (y_true == 1)) * ia_row).sum()\n        fp = ((pred_bin & (y_true == 0)) * ia_row).sum()\n        fn = (((1 - pred_bin) & (y_true == 1)) * ia_row).sum()\n\n        # Weighted precision, recall\n        prec = tp / (tp + fp + eps)\n        rec = tp / (tp + fn + eps)\n\n        # Weighted F1\n        f1 = 2.0 * prec * rec / (prec + rec + eps)\n\n        precisions.append(prec)\n        recalls.append(rec)\n        f1s.append(f1)\n        taus.append(tau)\n\n    curve_df = pd.DataFrame({\n        \"tau\": taus,\n        \"precision\": precisions,\n        \"recall\": recalls,\n        \"f1\": f1s,\n    })\n\n    # Best F-max and corresponding tau\n    best_idx = np.argmax(curve_df[\"f1\"].values)\n    best_f = float(curve_df.loc[best_idx, \"f1\"])\n    best_tau = float(curve_df.loc[best_idx, \"tau\"])\n\n    return best_f, best_tau, curve_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:58.867158Z","iopub.execute_input":"2025-12-10T11:58:58.867451Z","iopub.status.idle":"2025-12-10T11:58:58.877717Z","shell.execute_reply.started":"2025-12-10T11:58:58.867408Z","shell.execute_reply":"2025-12-10T11:58:58.877115Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"labels_per_protein = (train_df.groupby(\"protein_id\")[\"go_id\"]\n                      .apply(lambda x: set(x.dropna()))\n                      .apply(len))\nX.shape, Y.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:58:58.878522Z","iopub.execute_input":"2025-12-10T11:58:58.879095Z","iopub.status.idle":"2025-12-10T11:59:03.932832Z","shell.execute_reply.started":"2025-12-10T11:58:58.879078Z","shell.execute_reply":"2025-12-10T11:59:03.931991Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build a pseudo-stratification by number of labels (kept as-is)\nstrata = labels_per_protein.loc[protein_ids].values\nX_tr, X_va, y_tr, y_va = train_test_split(\n    X, Y, test_size=0.1, random_state=42, \n    stratify=np.clip(strata, 0, 10)\n)\n\n# Convert to tensors\nX_tr_t = torch.tensor(X_tr, dtype=torch.float32)\ny_tr_t = torch.tensor(y_tr, dtype=torch.float32)\nX_va_t = torch.tensor(X_va, dtype=torch.float32)\ny_va_t = torch.tensor(y_va, dtype=torch.float32)\n\n# Dataloaders\nbatch_size = 128\ntrain_ds = TensorDataset(X_tr_t, y_tr_t)\nval_ds = TensorDataset(X_va_t, y_va_t)\n\npin = device.type == \"cuda\"\ntrain_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, pin_memory=pin)\nval_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, pin_memory=pin)\n\n\nin_dim = X_tr.shape[1]\nout_dim = y_tr.shape[1] if y_tr.ndim == 2 else int(y_tr.max() + 1)\n\nclf = MLP(in_dim, out_dim).to(device)  # keep name \"clf\" for continuity\n\n# ---- Optimizer and loss (multi-label: BCEWithLogitsLoss) ----\noptimizer = torch.optim.Adam(clf.parameters(), lr=1e-3)\nscheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=1)\ncriterion = nn.BCEWithLogitsLoss()\n\n# AMP for speed on GPU\nuse_amp = device.type == \"cuda\"\nscaler = torch.cuda.amp.GradScaler(enabled=use_amp)\n\n# ---- Training loop (max_iter=30 equivalent to 30 epochs) ----\nepochs = 30\n\nfor epoch in range(1, epochs + 1):\n    clf.train()\n    for xb, yb in train_loader:\n        xb = xb.to(device, non_blocking=True)\n        yb = yb.to(device, non_blocking=True)\n\n        optimizer.zero_grad(set_to_none=True)\n        with torch.cuda.amp.autocast(enabled=use_amp):\n            logits = clf(xb)\n            loss = criterion(logits, yb) if y_tr.ndim == 2 else criterion(logits, yb.long())\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n\n    va_loss = evaluate(clf, val_loader)\n    print(f\"Epoch {epoch:02d}/{epochs} - val_loss: {va_loss:.5f}\")\n\n    # ---- Scheduler update ----\n    # ReduceLROnPlateau must be called AFTER validation\n    scheduler.step(va_loss)\n\n# ---- Validation predictions ----\nclf.eval()\nwith torch.no_grad():\n    va_logits = []\n    for xb, _ in val_loader:\n        xb = xb.to(device, non_blocking=True)\n        va_logits.append(clf(xb).detach().cpu())\n    va_logits = torch.cat(va_logits, dim=0)\n    va_logits = torch.sigmoid(va_logits)\n\ny_va_pred_proba_bce = va_logits.numpy()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:59:03.933636Z","iopub.execute_input":"2025-12-10T11:59:03.933895Z","iopub.status.idle":"2025-12-10T11:59:51.551394Z","shell.execute_reply.started":"2025-12-10T11:59:03.933874Z","shell.execute_reply":"2025-12-10T11:59:51.550594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fmax_bce, tau_bce, curve_df_bce = cafa_weighted_fmax(y_va, y_va_pred_proba_bce, term_weights)\n\nprint(fmax_bce, tau_bce)\n\ndef plot(curve_df):\n    import matplotlib.pyplot as plt\n    plt.figure(figsize=(10,6))\n    \n    plt.plot(curve_df[\"tau\"], curve_df[\"precision\"], label=\"Precision\")\n    plt.plot(curve_df[\"tau\"], curve_df[\"recall\"], label=\"Recall\")\n    plt.plot(curve_df[\"tau\"], curve_df[\"f1\"], label=\"F1\", linewidth=2)\n    \n    plt.xlabel(\"Tau (threshold)\")\n    plt.ylabel(\"Score\")\n    plt.title(\"CAFA Weighted Precision/Recall/F1 vs Tau\")\n    plt.grid(True, alpha=0.3)\n    plt.legend()\n    plt.show()\n\nplot(curve_df_bce)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T11:59:51.552324Z","iopub.execute_input":"2025-12-10T11:59:51.552615Z","iopub.status.idle":"2025-12-10T12:00:14.068789Z","shell.execute_reply.started":"2025-12-10T11:59:51.552590Z","shell.execute_reply":"2025-12-10T12:00:14.068103Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"EXPORT_PATH = \"cafa6_esm2_mlp_leaf1493_bundle.pt\"\n\nbundle = {\n    \"model_state_dict\": clf.state_dict(),\n    \"max_len\": MAX_LEN,\n    \"in_dim\": in_dim,\n    \"out_dim\": out_dim,\n    \"go_ids\": mlb.classes_,     # label order == Y columns\n    \"term_weights\": term_weights,\n    \"best_tau\": tau_bce,\n}\n\ntorch.save(bundle, EXPORT_PATH)\nprint(\"Saved bundle to\", EXPORT_PATH)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:00:14.069612Z","iopub.execute_input":"2025-12-10T12:00:14.069918Z","iopub.status.idle":"2025-12-10T12:00:14.085322Z","shell.execute_reply.started":"2025-12-10T12:00:14.069897Z","shell.execute_reply":"2025-12-10T12:00:14.084760Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}