{"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":14875579,"sourceType":"competition"},{"sourceId":13699671,"sourceType":"datasetVersion","datasetId":8714434},{"sourceId":14083674,"sourceType":"datasetVersion","datasetId":8960093},{"sourceId":276367826,"sourceType":"kernelVersion"},{"sourceId":285164560,"sourceType":"kernelVersion"},{"sourceId":285314858,"sourceType":"kernelVersion"},{"sourceId":676483,"sourceType":"modelInstanceVersion","modelInstanceId":512908,"modelId":527547}],"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:38:53.682533Z","iopub.execute_input":"2025-11-25T03:38:53.682866Z","iopub.status.idle":"2025-11-25T03:39:02.724799Z","shell.execute_reply.started":"2025-11-25T03:38:53.682842Z","shell.execute_reply":"2025-11-25T03:39:02.723948Z"}},"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\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 torch import nn\nfrom transformers import AutoTokenizer, AutoModel\nfrom einops import reduce\n\nimport csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:39:02.726272Z","iopub.execute_input":"2025-11-25T03:39:02.726503Z","iopub.status.idle":"2025-11-25T03:39:13.338064Z","shell.execute_reply.started":"2025-11-25T03:39:02.726479Z","shell.execute_reply":"2025-11-25T03:39:13.337229Z"}},"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\"            # GO:xxxx -> weight\nFP_OBO         = f\"{DATA_DIR}/Train/go-basic.obo\"      # optional: DAG analysis","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:39:22.7252Z","iopub.execute_input":"2025-11-25T03:39:22.726112Z","iopub.status.idle":"2025-11-25T03:39:22.730371Z","shell.execute_reply.started":"2025-11-25T03:39:22.726085Z","shell.execute_reply":"2025-11-25T03:39:22.729647Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Load bundle\nbundle = torch.load('/kaggle/input/cafa6-esm-dnn-training/cafa6_esm2_mlp_leaf1493_bundle.pt', map_location=device, weights_only=False)\nMAX_LEN     = bundle[\"max_len\"]\nin_dim      = bundle[\"in_dim\"]\nout_dim     = bundle[\"out_dim\"]\ngo_ids      = bundle[\"go_ids\"]\nterm_weights = bundle[\"term_weights\"]   # not used in submission, but kept\nbest_tau    = float(bundle[\"best_tau\"])\n\nprint(\"Loaded bundle:\")\nprint(\"  in_dim:\", in_dim, \"out_dim:\", out_dim)\nprint(\"  #GO terms:\", len(go_ids))\nprint(\"  best_tau:\", best_tau)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T04:32:35.502954Z","iopub.execute_input":"2025-11-25T04:32:35.503628Z","iopub.status.idle":"2025-11-25T04:32:35.527556Z","shell.execute_reply.started":"2025-11-25T04:32:35.503606Z","shell.execute_reply":"2025-11-25T04:32:35.526995Z"}},"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\ntest_seqs  = load_fasta(FP_TEST_FASTA)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:44:08.205832Z","iopub.execute_input":"2025-11-25T03:44:08.206526Z","iopub.status.idle":"2025-11-25T03:44:10.317014Z","shell.execute_reply.started":"2025-11-25T03:44:08.206493Z","shell.execute_reply":"2025-11-25T03:44:10.316217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = pd.DataFrame({\"protein_id\": list(test_seqs.keys()),\n                         \"sequence\":   list(test_seqs.values())})\nprotein_ids = test_df['protein_id'].sort_values().unique().tolist()\nprotein_ids[:10]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================\n# Cell 4: Embedding Backbone (ESM2 / ProtBERT / ProtT5)\n# =========================\n\n# Fast start; swap to a bigger checkpoint later for stronger features\nmodel_name = \"/kaggle/input/cafa6-esm2-model/esm2_t6_8m_UR50D\"\ntokenizer = AutoTokenizer.from_pretrained(model_name)\nbackbone  = AutoModel.from_pretrained(model_name).to(device)\n_ = backbone.eval()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:46:33.654808Z","iopub.execute_input":"2025-11-25T03:46:33.655139Z","iopub.status.idle":"2025-11-25T03:46:59.299151Z","shell.execute_reply.started":"2025-11-25T03:46:33.655118Z","shell.execute_reply":"2025-11-25T03:46:59.298334Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"MAX_LEN = 1024   # truncate long sequences for speed\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 = backbone(**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\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:46:59.300482Z","iopub.execute_input":"2025-11-25T03:46:59.301374Z","iopub.status.idle":"2025-11-25T03:46:59.307367Z","shell.execute_reply.started":"2025-11-25T03:46:59.301346Z","shell.execute_reply":"2025-11-25T03:46:59.306668Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"H = backbone.config.hidden_size\nX = np.zeros((len(protein_ids), H), dtype=np.float32)\nfor i, pid in enumerate(tqdm(protein_ids)):\n    X[i] = embed_sequence(test_seqs[pid])\nprint(f\"Embedding completed, shape: {X.shape}\")\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T03:48:38.43887Z","iopub.execute_input":"2025-11-25T03:48:38.439772Z","iopub.status.idle":"2025-11-25T04:27:37.202466Z","shell.execute_reply.started":"2025-11-25T03:48:38.439717Z","shell.execute_reply":"2025-11-25T04:27:37.201849Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# protein_ids = np.load('/kaggle/input/cafa6-protbert-embedding-test/protein_ids_test.npy')\n# X = np.load('/kaggle/input/cafa6-protbert-embedding-test/esm_testset.npy')\n# print(f\"Embedding completed, shape: {X.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- Torch setup ----\ndevice = 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)\n\nfrom torch.utils.data import TensorDataset, DataLoader\n\nbatch_size = 64\npin = device.type == \"cuda\"\nX_t = torch.tensor(X, dtype=torch.float32)\nds = TensorDataset(X_t)\nloader = DataLoader(ds, batch_size=batch_size, shuffle=False, pin_memory=pin)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T04:29:37.315534Z","iopub.execute_input":"2025-11-25T04:29:37.315861Z","iopub.status.idle":"2025-11-25T04:29:37.441098Z","shell.execute_reply.started":"2025-11-25T04:29:37.31584Z","shell.execute_reply":"2025-11-25T04:29:37.440445Z"}},"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\nin_dim = X_t.shape[1]\n\nclf = MLP(in_dim, out_dim).to(device)\nclf.load_state_dict(bundle[\"model_state_dict\"])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T04:32:41.280292Z","iopub.execute_input":"2025-11-25T04:32:41.280803Z","iopub.status.idle":"2025-11-25T04:32:41.302234Z","shell.execute_reply.started":"2025-11-25T04:32:41.280779Z","shell.execute_reply":"2025-11-25T04:32:41.301646Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_and_write_submission_streamed(\n    model,\n    loader,\n    protein_ids,           # list/array of length len(loader.dataset)\n    go_ids,                # list/array of length n_classes\n    out_path=\"/kaggle/working/submission.tsv\",\n    threshold=0.11,\n    apply_softmax=True     # set False if model(x) already returns probabilities\n):\n    \"\"\"\n    Stream inference and write rows where prob > threshold to a TSV file.\n    Memory footprint stays ~O(batch_size * n_classes) with no full-matrix materialization.\n    \"\"\"\n\n    # Safety checks\n    assert len(protein_ids) == len(loader.dataset), \"protein_ids length must match dataset length\"\n    device = next(model.parameters()).device\n    model.eval()\n\n    # Ensure loader is not shuffled so batch order matches protein_ids order\n    # (Set shuffle=False in its DataLoader)\n    row_base = 0  # global row offset (protein index)\n\n    with open(out_path, \"w\", newline=\"\") as f:\n        writer = csv.writer(f, delimiter=\"\\t\")\n        # No header per your format (header=False)\n\n        with torch.inference_mode():\n            for i, batch in enumerate(tqdm(loader)):\n                # Support (inputs,) or (inputs, labels, ...) batch structures\n                xb = batch[0] if isinstance(batch, (list, tuple)) else batch\n                xb = xb.to(device, non_blocking=True)\n\n                logits = model(xb)  # [B, C]\n                probs = torch.sigmoid(logits)\n\n                # Downcast on CPU to reduce memory while scanning threshold\n                probs_np = probs.to(torch.float16).cpu().numpy()  # shape [B, C]\n                B, C = probs_np.shape\n                # print((probs_np > threshold).sum())\n                # Find positions above threshold within the batch\n                bi, gi = np.where(probs_np > threshold)  # bi in [0..B-1], gi in [0..C-1]\n                if bi.size > 0:\n                    # Map local batch indices to global protein indices\n                    global_pi = (row_base + bi).astype(np.int64)\n\n                    # Vectorized fetch of values and ids\n                    # Note: converting to float32 string for stable formatting, keep as needed\n                    vals = probs_np[bi, gi].astype(np.float32)\n\n                    # Write rows: protein_id, go_id, prob\n                    # Avoid pandas here to keep peak memory minimal\n                    for p, g, v in zip(global_pi, gi, vals):\n                        writer.writerow([protein_ids[p], go_ids[g], f\"{v:.6f}\"])\n\n                # Free ASAP\n                del xb, logits, probs, probs_np\n                row_base += B\n\n                if i % 10 == 0:\n                    torch.cuda.empty_cache()\n\n    return out_path\n\n\nthreshold = 0.03\nsub_path = predict_and_write_submission_streamed(\n    clf,\n    loader,\n    protein_ids=protein_ids,\n    go_ids=go_ids,\n    out_path=\"submission.tsv\",\n    threshold=threshold,\n    apply_softmax=False   # set False if your model already outputs probabilities\n)\nprint(\"Written:\", sub_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T04:33:28.519425Z","iopub.execute_input":"2025-11-25T04:33:28.519918Z","iopub.status.idle":"2025-11-25T04:33:36.343702Z","shell.execute_reply.started":"2025-11-25T04:33:28.519895Z","shell.execute_reply":"2025-11-25T04:33:36.342981Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!wc -l submission.tsv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-25T04:33:40.747242Z","iopub.execute_input":"2025-11-25T04:33:40.747796Z","iopub.status.idle":"2025-11-25T04:33:40.946864Z","shell.execute_reply.started":"2025-11-25T04:33:40.747768Z","shell.execute_reply":"2025-11-25T04:33:40.946113Z"}},"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}]}