{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":13482672,"sourceType":"datasetVersion","datasetId":8559888,"isSourceIdPinned":false}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Dataset embeds","metadata":{}},{"cell_type":"code","source":"import kagglehub\n\npath = kagglehub.dataset_download(\"seddiktrk/cafa6-protein-embeddings-esm2\")\n\nprint(\"Path to dataset files:\", path)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:21.429553Z","iopub.execute_input":"2026-01-18T13:30:21.430190Z","iopub.status.idle":"2026-01-18T13:30:22.021688Z","shell.execute_reply.started":"2026-01-18T13:30:21.430159Z","shell.execute_reply":"2026-01-18T13:30:22.021045Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nprint(os.listdir(\"/kaggle/input\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:22.023129Z","iopub.execute_input":"2026-01-18T13:30:22.023504Z","iopub.status.idle":"2026-01-18T13:30:22.027377Z","shell.execute_reply.started":"2026-01-18T13:30:22.023480Z","shell.execute_reply":"2026-01-18T13:30:22.026806Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, re, json\nimport numpy as np\nimport pandas as pd\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.amp import autocast, GradScaler","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:22.028086Z","iopub.execute_input":"2026-01-18T13:30:22.028377Z","iopub.status.idle":"2026-01-18T13:30:23.347348Z","shell.execute_reply.started":"2026-01-18T13:30:22.028349Z","shell.execute_reply":"2026-01-18T13:30:23.346676Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"COMP_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\nTRAIN_DIR = os.path.join(COMP_DIR, \"Train\")\nTEST_DIR  = os.path.join(COMP_DIR, \"Test\")\n\nTRAIN_TERMS = os.path.join(TRAIN_DIR, \"train_terms.tsv\")\nTRAIN_FASTA = os.path.join(TRAIN_DIR, \"train_sequences.fasta\")\nTEST_FASTA  = os.path.join(TEST_DIR,  \"testsuperset.fasta\")\nOBO_PATH    = os.path.join(TRAIN_DIR, \"go-basic.obo\")\n\nEMB_DIR = \"/kaggle/input/cafa6-protein-embeddings-esm2\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:23.350133Z","iopub.execute_input":"2026-01-18T13:30:23.350658Z","iopub.status.idle":"2026-01-18T13:30:23.355667Z","shell.execute_reply.started":"2026-01-18T13:30:23.350630Z","shell.execute_reply":"2026-01-18T13:30:23.354954Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"FASTA loader + ID normalization","metadata":{}},{"cell_type":"code","source":"def read_fasta_to_df(path: str) -> pd.DataFrame:\n    ids, seqs = [], []\n    cur_id, cur_seq = None, []\n    with open(path, \"r\", encoding=\"utf-8\") as f:\n        for raw in f:\n            line = raw.strip()\n            if not line: \n                continue\n            if line.startswith(\">\"):\n                if cur_id is not None:\n                    ids.append(cur_id); seqs.append(\"\".join(cur_seq))\n                cur_id = line[1:].split()[0]\n                cur_seq = []\n            else:\n                cur_seq.append(line)\n    if cur_id is not None:\n        ids.append(cur_id); seqs.append(\"\".join(cur_seq))\n    return pd.DataFrame({\"protein_id\": ids, \"sequence\": seqs})\n\ndef normalize_protein_id(x: str) -> str:\n    x = str(x).strip().split()[0]\n    if \"|\" in x:\n        parts = x.split(\"|\")\n        if len(parts) >= 2:\n            x = parts[1]\n        else:\n            x = parts[-1]\n    if \".\" in x:\n        left, right = x.rsplit(\".\", 1)\n        if right.isdigit():\n            x = left\n    return x\n\ntrain_seq = read_fasta_to_df(TRAIN_FASTA)\ntest_seq  = read_fasta_to_df(TEST_FASTA)\n\ntrain_seq[\"protein_id\"] = train_seq[\"protein_id\"].map(normalize_protein_id)\ntest_seq[\"protein_id\"]  = test_seq[\"protein_id\"].map(normalize_protein_id)\n\ntrain_terms = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\ntrain_terms = train_terms.rename(columns={\"EntryID\":\"protein_id\",\"term\":\"go_id\"})\ntrain_terms[\"protein_id\"] = train_terms[\"protein_id\"].map(normalize_protein_id)\n\nprint(train_seq.shape, test_seq.shape, train_terms.shape)\nprint(train_terms[\"aspect\"].value_counts())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:23.356702Z","iopub.execute_input":"2026-01-18T13:30:23.357004Z","iopub.status.idle":"2026-01-18T13:30:26.239588Z","shell.execute_reply.started":"2026-01-18T13:30:23.356980Z","shell.execute_reply":"2026-01-18T13:30:26.238785Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Split train/val (no duplication)","metadata":{}},{"cell_type":"code","source":"SEED = 42\nVAL_FRAC = 0.10\n\nseq_to_rep = train_seq.groupby(\"sequence\")[\"protein_id\"].min().rename(\"rep_id\").reset_index()\ntrain_rep = train_seq.merge(seq_to_rep, on=\"sequence\", how=\"left\")\n\nrep_ids = train_rep[\"rep_id\"].drop_duplicates().values\nrng = np.random.default_rng(SEED); rng.shuffle(rep_ids)\n\nn_val = int(len(rep_ids) * VAL_FRAC)\nval_rep = set(rep_ids[:n_val])\ntrain_rep_set = set(rep_ids[n_val:])\n\nval_ids   = train_rep.loc[train_rep[\"rep_id\"].isin(val_rep), \"protein_id\"].unique().tolist()\ntrain_ids = train_rep.loc[train_rep[\"rep_id\"].isin(train_rep_set), \"protein_id\"].unique().tolist()\n\nprint(\"train:\", len(train_ids), \"val:\", len(val_ids), \"overlap:\", len(set(train_ids)&set(val_ids)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:26.240628Z","iopub.execute_input":"2026-01-18T13:30:26.240904Z","iopub.status.idle":"2026-01-18T13:30:30.016031Z","shell.execute_reply.started":"2026-01-18T13:30:26.240872Z","shell.execute_reply":"2026-01-18T13:30:30.015286Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Load embeddings","metadata":{}},{"cell_type":"code","source":"def find_file(patterns):\n    files = []\n    for root, _, fnames in os.walk(EMB_DIR):\n        for f in fnames:\n            files.append(os.path.join(root, f))\n    for p in patterns:\n        cand = [x for x in files if p in os.path.basename(x).lower()]\n        if cand:\n            return cand[0]\n    return None\n\nprint(\"EMB_DIR listing (top):\", os.listdir(EMB_DIR)[:20])\n\ntrain_emb_path = find_file([\"train\", \"emb\", \".npy\"]) or find_file([\"train\", \".parquet\"]) or find_file([\"train\", \".npz\"])\ntest_emb_path  = find_file([\"test\",  \"emb\", \".npy\"]) or find_file([\"test\",  \".parquet\"]) or find_file([\"test\",  \".npz\"])\nid_path_train  = find_file([\"train\", \"id\"]) or find_file([\"train\", \"protein\"])\nid_path_test   = find_file([\"test\",  \"id\"]) or find_file([\"test\",  \"protein\"])\n\nprint(\"train_emb_path:\", train_emb_path)\nprint(\"test_emb_path :\", test_emb_path)\nprint(\"id_path_train :\", id_path_train)\nprint(\"id_path_test  :\", id_path_test)\n\ndef load_embeddings(emb_path, id_path=None):\n    if emb_path.endswith(\".parquet\"):\n        df = pd.read_parquet(emb_path)\n        if \"protein_id\" in df.columns and \"embedding\" in df.columns:\n            ids = df[\"protein_id\"].astype(str).tolist()\n            X = np.vstack(df[\"embedding\"].to_numpy())\n            return ids, X\n        elif \"protein_id\" in df.columns:\n            ids = df[\"protein_id\"].astype(str).tolist()\n            X = df.drop(columns=[\"protein_id\"]).to_numpy()\n            return ids, X\n        else:\n            raise ValueError(\"Parquet embeddings missing protein_id column.\")\n    elif emb_path.endswith(\".npz\"):\n        z = np.load(emb_path, allow_pickle=True)\n        keyX = \"embeddings\" if \"embeddings\" in z else (\"X\" if \"X\" in z else None)\n        keyI = \"protein_id\" if \"protein_id\" in z else (\"ids\" if \"ids\" in z else None)\n        if keyX is None or keyI is None:\n            raise ValueError(f\"npz keys: {list(z.keys())}\")\n        ids = [str(x) for x in z[keyI]]\n        X = z[keyX]\n        return ids, X\n    elif emb_path.endswith(\".npy\"):\n        X = np.load(emb_path, mmap_mode=\"r\")\n        if id_path is None:\n            raise ValueError(\"Need an id file to align .npy embeddings.\")\n        if id_path.endswith(\".csv\") or id_path.endswith(\".tsv\"):\n            ids = pd.read_csv(id_path, header=None, sep=\"\\t\" if id_path.endswith(\".tsv\") else \",\")[0].astype(str).tolist()\n        else:\n            ids = [str(x) for x in np.load(id_path, allow_pickle=True)]\n        return ids, X\n    else:\n        raise ValueError(\"Unknown embedding format: \" + emb_path)\n\ntrain_emb_ids, X_train_all = load_embeddings(train_emb_path, id_path_train)\ntest_emb_ids,  X_test_all  = load_embeddings(test_emb_path,  id_path_test)\n\ntrain_emb_ids = [normalize_protein_id(x) for x in train_emb_ids]\ntest_emb_ids  = [normalize_protein_id(x) for x in test_emb_ids]\n\nprint(\"X_train_all:\", np.shape(X_train_all), \"X_test_all:\", np.shape(X_test_all))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:30.016959Z","iopub.execute_input":"2026-01-18T13:30:30.017247Z","iopub.status.idle":"2026-01-18T13:30:30.390575Z","shell.execute_reply.started":"2026-01-18T13:30:30.017207Z","shell.execute_reply":"2026-01-18T13:30:30.389729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if len(test_emb_ids) != X_test_all.shape[0]:\n    n = min(len(test_emb_ids), X_test_all.shape[0])\n    print(\"Aligning test ids/embeddings to length:\", n)\n    test_emb_ids = test_emb_ids[:n]\n    X_test_all = X_test_all[:n]\n\ntest_id_to_row = {pid:i for i, pid in enumerate(test_emb_ids)}\n\nprint(\"after align -> ids:\", len(test_emb_ids), \"X:\", X_test_all.shape[0])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:30.391675Z","iopub.execute_input":"2026-01-18T13:30:30.391953Z","iopub.status.idle":"2026-01-18T13:30:30.449847Z","shell.execute_reply.started":"2026-01-18T13:30:30.391931Z","shell.execute_reply":"2026-01-18T13:30:30.449000Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if len(test_emb_ids) == X_test_all.shape[0] + 1 and test_emb_ids[-1] == \"A4QQE0\":\n    test_emb_ids = test_emb_ids[:-1]\n    # X_test_all unchanged\ntest_id_to_row = {pid:i for i, pid in enumerate(test_emb_ids)}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:30.450795Z","iopub.execute_input":"2026-01-18T13:30:30.451133Z","iopub.status.idle":"2026-01-18T13:30:30.521046Z","shell.execute_reply.started":"2026-01-18T13:30:30.451109Z","shell.execute_reply":"2026-01-18T13:30:30.520203Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_split_embeddings(pids, X_all, id_to_row):\n    rows = []\n    kept = []\n    for pid in pids:\n        r = id_to_row.get(pid, None)\n        if r is None:\n            continue\n        rows.append(r)\n        kept.append(pid)\n    X = np.asarray(X_all[rows], dtype=np.float32)\n    return kept, X","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:30.523521Z","iopub.execute_input":"2026-01-18T13:30:30.523776Z","iopub.status.idle":"2026-01-18T13:30:30.528385Z","shell.execute_reply.started":"2026-01-18T13:30:30.523752Z","shell.execute_reply":"2026-01-18T13:30:30.527556Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_pids_aligned, X_test = make_split_embeddings(\n    test_seq[\"protein_id\"].tolist(), X_test_all, test_id_to_row\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:30:30.529215Z","iopub.execute_input":"2026-01-18T13:30:30.529897Z","iopub.status.idle":"2026-01-18T13:31:05.624621Z","shell.execute_reply.started":"2026-01-18T13:30:30.529846Z","shell.execute_reply":"2026-01-18T13:31:05.623687Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"aligned test proteins:\", len(test_pids_aligned), \"expected:\", len(test_seq))\nprint(\"missing:\", len(test_seq) - len(test_pids_aligned))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:05.625922Z","iopub.execute_input":"2026-01-18T13:31:05.626367Z","iopub.status.idle":"2026-01-18T13:31:05.630903Z","shell.execute_reply.started":"2026-01-18T13:31:05.626328Z","shell.execute_reply":"2026-01-18T13:31:05.630088Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"missing = list(set(test_seq[\"protein_id\"]) - set(test_pids_aligned))\nmissing[:10], len(missing)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:05.631828Z","iopub.execute_input":"2026-01-18T13:31:05.632525Z","iopub.status.idle":"2026-01-18T13:31:05.702353Z","shell.execute_reply.started":"2026-01-18T13:31:05.632493Z","shell.execute_reply":"2026-01-18T13:31:05.701561Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**clean fallback for just A4QQE0**","metadata":{}},{"cell_type":"code","source":"top_fallback = {}\nfor asp in [\"BP\",\"MF\",\"CC\"]:\n    s = train_terms[(train_terms[\"protein_id\"].isin(train_ids)) & (train_terms[\"aspect\"]==asp)][\"go_id\"].value_counts()\n    top_fallback[asp] = s.head(200).index.tolist()\n    print(asp, \"fallback terms:\", len(top_fallback[asp]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:05.703370Z","iopub.execute_input":"2026-01-18T13:31:05.704246Z","iopub.status.idle":"2026-01-18T13:31:05.925184Z","shell.execute_reply.started":"2026-01-18T13:31:05.704222Z","shell.execute_reply":"2026-01-18T13:31:05.924540Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"debugging BP fallback terms: 0\nMF fallback terms: 0\nCC fallback terms: 0","metadata":{}},{"cell_type":"code","source":"print(\"len(train_ids):\", len(train_ids))\nprint(\"train_terms proteins:\", train_terms[\"protein_id\"].nunique())\n\nmatches = train_terms[\"protein_id\"].isin(set(train_ids)).sum()\nprint(\"matches rows:\", matches)\n\nprint(\"example train_ids:\", train_ids[:5])\nprint(\"example train_terms ids:\", train_terms[\"protein_id\"].head(5).tolist())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:05.926083Z","iopub.execute_input":"2026-01-18T13:31:05.926360Z","iopub.status.idle":"2026-01-18T13:31:06.008159Z","shell.execute_reply.started":"2026-01-18T13:31:05.926337Z","shell.execute_reply":"2026-01-18T13:31:06.007578Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**fallback from train_terms**\n","metadata":{}},{"cell_type":"code","source":"top_fallback = {}\nfor asp in [\"BP\",\"MF\",\"CC\"]:\n    s = train_terms[train_terms[\"aspect\"]==asp][\"go_id\"].value_counts()\n    top_fallback[asp] = s.head(200).index.tolist()\n    print(asp, \"fallback terms:\", len(top_fallback[asp]))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:06.009155Z","iopub.execute_input":"2026-01-18T13:31:06.009527Z","iopub.status.idle":"2026-01-18T13:31:06.130055Z","shell.execute_reply.started":"2026-01-18T13:31:06.009495Z","shell.execute_reply":"2026-01-18T13:31:06.129329Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(train_terms[\"aspect\"].value_counts().head(20))\nprint(\"unique aspects:\", sorted(train_terms[\"aspect\"].astype(str).unique())[:20])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:06.130887Z","iopub.execute_input":"2026-01-18T13:31:06.131120Z","iopub.status.idle":"2026-01-18T13:31:06.188668Z","shell.execute_reply.started":"2026-01-18T13:31:06.131100Z","shell.execute_reply":"2026-01-18T13:31:06.187819Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms[\"aspect\"] = train_terms[\"aspect\"].astype(str).str.strip().str.upper()\n\naspect_map = {\n    \"BPO\": \"BP\", \"BP\": \"BP\", \"P\": \"BP\", \"BIOLOGICAL_PROCESS\": \"BP\",\n    \"MFO\": \"MF\", \"MF\": \"MF\", \"F\": \"MF\", \"MOLECULAR_FUNCTION\": \"MF\",\n    \"CCO\": \"CC\", \"CC\": \"CC\", \"C\": \"CC\", \"CELLULAR_COMPONENT\": \"CC\",\n}\n\ntrain_terms[\"asp_norm\"] = train_terms[\"aspect\"].map(aspect_map)\nprint(train_terms[\"asp_norm\"].value_counts(dropna=False))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:06.189702Z","iopub.execute_input":"2026-01-18T13:31:06.190003Z","iopub.status.idle":"2026-01-18T13:31:06.424052Z","shell.execute_reply.started":"2026-01-18T13:31:06.189980Z","shell.execute_reply":"2026-01-18T13:31:06.423293Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"top_fallback = {}\nfor asp in [\"BP\",\"MF\",\"CC\"]:\n    s = train_terms[(train_terms[\"asp_norm\"]==asp)][\"go_id\"].value_counts()\n    top_fallback[asp] = s.head(200).index.tolist()\n    print(asp, \"fallback terms:\", len(top_fallback[asp]), \"top1:\", top_fallback[asp][0] if top_fallback[asp] else None)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:06.425135Z","iopub.execute_input":"2026-01-18T13:31:06.425431Z","iopub.status.idle":"2026-01-18T13:31:06.636648Z","shell.execute_reply.started":"2026-01-18T13:31:06.425405Z","shell.execute_reply":"2026-01-18T13:31:06.636076Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Label vocab per aspect + multi-hot targets","metadata":{}},{"cell_type":"code","source":"ASPECTS = [\"BP\", \"MF\", \"CC\"]\n\nterms_train_split = train_terms[train_terms[\"protein_id\"].isin(train_ids)].copy()\n\ngo_vocab = {}\ngo2i = {}\n\nfor asp in ASPECTS:\n    vocab = sorted(terms_train_split.loc[terms_train_split[\"asp_norm\"]==asp, \"go_id\"].unique())\n    go_vocab[asp] = vocab\n    go2i[asp] = {go:i for i,go in enumerate(vocab)}\n    print(asp, \"labels:\", len(vocab))\n\ndef pid2labels_for_aspect(pids_set, asp):\n    df = train_terms[(train_terms[\"protein_id\"].isin(pids_set)) & (train_terms[\"asp_norm\"]==asp)][[\"protein_id\",\"go_id\"]].copy()\n    df = df[df[\"go_id\"].isin(go2i[asp])]\n    df[\"y\"] = df[\"go_id\"].map(go2i[asp]).astype(int)\n    return df.groupby(\"protein_id\")[\"y\"].apply(list).to_dict()\n\ntrain_pid2labels = {asp: pid2labels_for_aspect(set(train_ids), asp) for asp in ASPECTS}\nval_pid2labels   = {asp: pid2labels_for_aspect(set(val_ids),   asp) for asp in ASPECTS}\n\nprint(\"example BP labels for a train protein:\", next(iter(train_pid2labels[\"BP\"].items())))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:06.637581Z","iopub.execute_input":"2026-01-18T13:31:06.637820Z","iopub.status.idle":"2026-01-18T13:31:10.360048Z","shell.execute_reply.started":"2026-01-18T13:31:06.637799Z","shell.execute_reply":"2026-01-18T13:31:10.359388Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Align embeddings to split IDs","metadata":{}},{"cell_type":"code","source":"train_id_to_row = {pid:i for i,pid in enumerate(train_emb_ids)}\ntest_id_to_row  = {pid:i for i,pid in enumerate(test_emb_ids)}\n\ndef make_split_embeddings(pids, X_all, id_to_row):\n    rows = [id_to_row[pid] for pid in pids if pid in id_to_row]\n    kept = [pid for pid in pids if pid in id_to_row]\n    X = np.asarray(X_all[rows], dtype=np.float32)\n    return kept, X\n\ntrain_pids_aligned, X_train = make_split_embeddings(train_ids, X_train_all, train_id_to_row)\nval_pids_aligned,   X_val   = make_split_embeddings(val_ids,   X_train_all, train_id_to_row)\ntest_pids_aligned,  X_test  = make_split_embeddings(test_seq[\"protein_id\"].tolist(), X_test_all, test_id_to_row)\n\nprint(\"aligned train/val/test:\", len(train_pids_aligned), len(val_pids_aligned), len(test_pids_aligned))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:10.361010Z","iopub.execute_input":"2026-01-18T13:31:10.361264Z","iopub.status.idle":"2026-01-18T13:31:11.229381Z","shell.execute_reply.started":"2026-01-18T13:31:10.361241Z","shell.execute_reply":"2026-01-18T13:31:11.228569Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Train MLP heads with logits-stable version using binary_cross_entropy_with_logits","metadata":{}},{"cell_type":"code","source":"import torch.nn.functional as F\n\nclass ASLStable(nn.Module):\n    def __init__(self, gamma_pos=0.0, gamma_neg=2.0):\n        super().__init__()\n        self.gp = gamma_pos\n        self.gn = gamma_neg\n\n    def forward(self, logits, targets):\n        bce = F.binary_cross_entropy_with_logits(logits, targets, reduction=\"none\")\n\n        p = torch.sigmoid(logits)\n        pt = p*targets + (1-p)*(1-targets)\n\n        w = (1-pt).pow(self.gp*targets + self.gn*(1-targets))\n        return (w*bce).mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.230363Z","iopub.execute_input":"2026-01-18T13:31:11.230607Z","iopub.status.idle":"2026-01-18T13:31:11.235779Z","shell.execute_reply.started":"2026-01-18T13:31:11.230585Z","shell.execute_reply":"2026-01-18T13:31:11.235205Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"loss_fn = ASLStable(gamma_pos=0.0, gamma_neg=2.0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.236562Z","iopub.execute_input":"2026-01-18T13:31:11.236781Z","iopub.status.idle":"2026-01-18T13:31:11.254318Z","shell.execute_reply.started":"2026-01-18T13:31:11.236760Z","shell.execute_reply":"2026-01-18T13:31:11.253592Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Training each aspect**","metadata":{}},{"cell_type":"code","source":"class EmbDataset(Dataset):\n    def __init__(self, X, pids, pid2labels):\n        self.X = X\n        self.pids = pids\n        self.pid2labels = pid2labels\n\n    def __len__(self):\n        return len(self.pids)\n\n    def __getitem__(self, i):\n        pid = self.pids[i]\n        labs = self.pid2labels.get(pid, [])\n        return self.X[i], labs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.255347Z","iopub.execute_input":"2026-01-18T13:31:11.255567Z","iopub.status.idle":"2026-01-18T13:31:11.269058Z","shell.execute_reply.started":"2026-01-18T13:31:11.255548Z","shell.execute_reply":"2026-01-18T13:31:11.268433Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def collate_emb(batch, n_labels):\n    xs, labs_list = zip(*batch)\n    Xb = torch.tensor(np.stack(xs), dtype=torch.float32)\n    Yb = torch.zeros((len(xs), n_labels), dtype=torch.float32)\n    for i, labs in enumerate(labs_list):\n        if labs:\n            Yb[i, labs] = 1.0\n    return Xb, Yb","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.269729Z","iopub.execute_input":"2026-01-18T13:31:11.269945Z","iopub.status.idle":"2026-01-18T13:31:11.284666Z","shell.execute_reply.started":"2026-01-18T13:31:11.269925Z","shell.execute_reply":"2026-01-18T13:31:11.284035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ASLStable(nn.Module):\n    def __init__(self, gamma_pos=0.0, gamma_neg=2.0):\n        super().__init__()\n        self.gp = gamma_pos\n        self.gn = gamma_neg\n    def forward(self, logits, targets):\n        bce = F.binary_cross_entropy_with_logits(logits, targets, reduction=\"none\")\n        p = torch.sigmoid(logits)\n        pt = p*targets + (1-p)*(1-targets)\n        w = (1-pt).pow(self.gp*targets + self.gn*(1-targets))\n        return (w*bce).mean()\n\nclass EmbDataset(Dataset):\n    def __init__(self, X, pids, pid2labels):\n        self.X = X\n        self.pids = pids\n        self.pid2labels = pid2labels\n    def __len__(self): return len(self.pids)\n    def __getitem__(self, i):\n        pid = self.pids[i]\n        labs = self.pid2labels.get(pid, [])\n        return self.X[i], labs\n\ndef collate_emb(batch, n_labels):\n    xs, labs_list = zip(*batch)\n    Xb = torch.tensor(np.stack(xs), dtype=torch.float32)\n    Yb = torch.zeros((len(xs), n_labels), dtype=torch.float32)\n    for i, labs in enumerate(labs_list):\n        if labs:\n            Yb[i, labs] = 1.0\n    return Xb, Yb\n\nclass MLP(nn.Module):\n    def __init__(self, d_in, d_hid, n_out):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(d_in, d_hid),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(d_hid, n_out),\n        )\n    def forward(self, x): return self.net(x)\n\nD = X_train.shape[1]\nprint(\"embedding dim:\", D)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.285491Z","iopub.execute_input":"2026-01-18T13:31:11.285727Z","iopub.status.idle":"2026-01-18T13:31:11.300880Z","shell.execute_reply.started":"2026-01-18T13:31:11.285694Z","shell.execute_reply":"2026-01-18T13:31:11.300308Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SEED = 42\nVAL_FRAC = 0.10\n\nEPOCHS = 6\nBATCH = 1024\nLR = 1e-3\nNUM_WORKERS = 2\n\nTOPK_SUB = 200\nDEFAULT_THR = 0.20\nOUT_PATH = \"/kaggle/working/submission.tsv\"\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"device:\", device)\n\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif device.type == \"cuda\":\n    torch.cuda.manual_seed_all(SEED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.301760Z","iopub.execute_input":"2026-01-18T13:31:11.302039Z","iopub.status.idle":"2026-01-18T13:31:11.418868Z","shell.execute_reply.started":"2026-01-18T13:31:11.302010Z","shell.execute_reply":"2026-01-18T13:31:11.418245Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.amp import GradScaler\n\nscaler = GradScaler(\"cuda\", enabled=(device.type == \"cuda\"))\nprint(\"scaler enabled:\", scaler.is_enabled())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.422319Z","iopub.execute_input":"2026-01-18T13:31:11.422836Z","iopub.status.idle":"2026-01-18T13:31:11.430559Z","shell.execute_reply.started":"2026-01-18T13:31:11.422811Z","shell.execute_reply":"2026-01-18T13:31:11.429820Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"models = {}\nfor asp in ASPECTS:\n    nlab = len(go_vocab[asp])\n\n    train_ds = EmbDataset(X_train, train_pids_aligned, train_pid2labels[asp])\n    val_ds   = EmbDataset(X_val,   val_pids_aligned,   val_pid2labels[asp])\n\n    train_loader = DataLoader(\n        train_ds, batch_size=BATCH, shuffle=True,\n        num_workers=NUM_WORKERS, pin_memory=True,\n        collate_fn=lambda b, nlab=nlab: collate_emb(b, nlab)\n    )\n    val_loader = DataLoader(\n        val_ds, batch_size=BATCH, shuffle=False,\n        num_workers=NUM_WORKERS, pin_memory=True,\n        collate_fn=lambda b, nlab=nlab: collate_emb(b, nlab)\n    )\n\n    model = MLP(D, 2048, nlab).to(device)\n    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)\n\n    # IMPORTANT: use the stable loss you created\n    loss_fn = ASLStable(gamma_pos=0.0, gamma_neg=2.0)\n\n    for ep in range(1, EPOCHS + 1):\n        model.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            opt.zero_grad(set_to_none=True)\n            with autocast(\"cuda\", enabled=(device.type == \"cuda\")):\n                logits = model(Xb)\n                loss = loss_fn(logits, Yb)\n\n            scaler.scale(loss).backward()\n            scaler.step(opt)\n            scaler.update()\n\n        print(f\"{asp} epoch {ep} done\")\n\n    models[asp] = model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:31:11.431490Z","iopub.execute_input":"2026-01-18T13:31:11.431718Z","iopub.status.idle":"2026-01-18T13:32:09.018238Z","shell.execute_reply.started":"2026-01-18T13:31:11.431698Z","shell.execute_reply":"2026-01-18T13:32:09.017357Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Validation thresholds (per aspect) + NaN check","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport torch\n\ndef build_multihot(pids, pid2labels, n_labels):\n    Y = np.zeros((len(pids), n_labels), dtype=np.uint8)\n    for i, pid in enumerate(pids):\n        labs = pid2labels.get(pid, [])\n        if labs:\n            Y[i, labs] = 1\n    return Y\n\ndef fmax_safe(Y_true, Y_score, thresholds=np.linspace(0.01, 0.99, 99)):\n    best_f, best_t = -1.0, 0.20\n    for t in thresholds:\n        Y_pred = (Y_score >= t).astype(np.uint8)\n        tp = (Y_pred & Y_true).sum()\n        fp = (Y_pred & (1 - Y_true)).sum()\n        fn = ((1 - Y_pred) & Y_true).sum()\n        p = tp / (tp + fp + 1e-9)\n        r = tp / (tp + fn + 1e-9)\n        f = 2 * p * r / (p + r + 1e-9)\n        if f > best_f:\n            best_f, best_t = float(f), float(t)\n    return best_f, best_t\n\nthr_aspect = {}\n\nXv = torch.tensor(X_val, dtype=torch.float32)\nfor asp in ASPECTS:\n    model = models[asp].to(device).eval()\n    nlab = len(go_vocab[asp])\n    Y_true = build_multihot(val_pids_aligned, val_pid2labels[asp], nlab)\n\n    outs = []\n    bs = 4096\n    with torch.no_grad():\n        for i in range(0, Xv.shape[0], bs):\n            xb = Xv[i:i+bs].to(device, non_blocking=True)\n            prob = torch.sigmoid(model(xb)).float().cpu().numpy()\n            outs.append(prob)\n    Y_score = np.vstack(outs).astype(np.float32)\n\n    print(asp, \"Y_score:\", \"min\", float(np.nanmin(Y_score)), \"mean\", float(np.nanmean(Y_score)), \"max\", float(np.nanmax(Y_score)))\n    if np.isnan(Y_score).any():\n        raise ValueError(f\"{asp}: NaNs detected in Y_score (do NOT proceed)\")\n\n    best_f, thr = fmax_safe(Y_true, Y_score)\n    thr_aspect[asp] = thr\n    print(asp, \"best_f:\", best_f, \"thr:\", thr)\n\nthr_aspect","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:32:09.019635Z","iopub.execute_input":"2026-01-18T13:32:09.020262Z","iopub.status.idle":"2026-01-18T13:33:30.132831Z","shell.execute_reply.started":"2026-01-18T13:32:09.020229Z","shell.execute_reply":"2026-01-18T13:33:30.132184Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Predict TEST + write submission.tsv ","metadata":{}},{"cell_type":"code","source":"Xt = torch.tensor(X_test, dtype=torch.float32)\n\ntest_pid_to_row = {pid:i for i,pid in enumerate(test_pids_aligned)}\nmissing_ids = list(set(test_seq[\"protein_id\"]) - set(test_pids_aligned))\n\nwritten = 0\nwith open(OUT_PATH, \"w\", encoding=\"utf-8\") as f:\n    for asp in ASPECTS:\n        model = models[asp].eval()\n        nlab = len(go_vocab[asp])\n        i2go = go_vocab[asp]\n        thr = float(thr_aspect.get(asp, DEFAULT_THR))\n\n        bs = 4096\n        with torch.no_grad():\n            for i in range(0, Xt.shape[0], bs):\n                xb = Xt[i:i+bs].to(device, non_blocking=True)\n                prob = torch.sigmoid(model(xb)).float().cpu()  # [B, nlab]\n\n                k = min(TOPK_SUB, nlab)\n                sc, idx = torch.topk(prob, k=k, dim=1)\n                sc = sc.numpy()\n                idx = idx.numpy()\n\n                for r in range(sc.shape[0]):\n                    pid = test_pids_aligned[i + r]\n                    for j in range(sc.shape[1]):\n                        s = float(sc[r, j])\n                        if s < thr:\n                            continue\n                        go = i2go[int(idx[r, j])]\n                        f.write(f\"{pid}\\t{go}\\t{s:.6f}\\n\")\n                        written += 1\n\n    # fallback for missing protein (A4QQE0)\n    for pid in missing_ids:\n        for asp in ASPECTS:\n            for go in top_fallback[asp]:\n                f.write(f\"{pid}\\t{go}\\t{DEFAULT_THR:.6f}\\n\")\n                written += 1\n\nprint(\"Wrote:\", OUT_PATH, \"| rows:\", written, \"| missing fallback:\", missing_ids[:5])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:33:30.133941Z","iopub.execute_input":"2026-01-18T13:33:30.134198Z","iopub.status.idle":"2026-01-18T13:35:25.752138Z","shell.execute_reply.started":"2026-01-18T13:33:30.134176Z","shell.execute_reply":"2026-01-18T13:35:25.751220Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv(OUT_PATH, sep=\"\\t\", header=None, names=[\"protein_id\",\"go_id\",\"score\"])\nprint(\"rows:\", len(df))\nprint(\"unique proteins:\", df[\"protein_id\"].nunique(), \"expected test:\", len(test_seq))\nprint(\"unique GO:\", df[\"go_id\"].nunique())\nprint(\"avg rows/protein:\", df.groupby(\"protein_id\").size().mean())\nprint(df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:25.753428Z","iopub.execute_input":"2026-01-18T13:35:25.753754Z","iopub.status.idle":"2026-01-18T13:35:27.220532Z","shell.execute_reply.started":"2026-01-18T13:35:25.753730Z","shell.execute_reply":"2026-01-18T13:35:27.219802Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv(OUT_PATH, sep=\"\\t\", header=None, names=[\"protein_id\",\"go_id\",\"score\"])\n\nprint(\"nan scores:\", df[\"score\"].isna().sum())\nprint(\"score min/max:\", df[\"score\"].min(), df[\"score\"].max())\nprint(\"bad GO format:\", (~df[\"go_id\"].str.match(r\"^GO:\\d{7}$\")).sum())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:27.221538Z","iopub.execute_input":"2026-01-18T13:35:27.221789Z","iopub.status.idle":"2026-01-18T13:35:29.075875Z","shell.execute_reply.started":"2026-01-18T13:35:27.221766Z","shell.execute_reply":"2026-01-18T13:35:29.075173Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nprint(os.listdir(\"/kaggle/input/cafa-6-protein-function-prediction\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:29.076936Z","iopub.execute_input":"2026-01-18T13:35:29.077572Z","iopub.status.idle":"2026-01-18T13:35:29.083041Z","shell.execute_reply.started":"2026-01-18T13:35:29.077543Z","shell.execute_reply":"2026-01-18T13:35:29.082359Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nSAMPLE_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\"\n\nsamp = pd.read_csv(SAMPLE_PATH, header=None, names=[\"protein_id\",\"go_id\",\"score\"])\nprint(samp.head())\nprint(\"shape:\", samp.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:29.084090Z","iopub.execute_input":"2026-01-18T13:35:29.084489Z","iopub.status.idle":"2026-01-18T13:35:29.111344Z","shell.execute_reply.started":"2026-01-18T13:35:29.084456Z","shell.execute_reply":"2026-01-18T13:35:29.110785Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SUB_PATH = \"/kaggle/working/submission.tsv\"\nsub = pd.read_csv(SUB_PATH, sep=\"\\t\", header=None, names=[\"protein_id\",\"go_id\",\"score\"])\n\nsub_ids = set(sub[\"protein_id\"].astype(str))\nsamp_ids = set(samp[\"protein_id\"].astype(str))\n\nprint(\"submission unique ids:\", len(sub_ids))\nprint(\"sample unique ids    :\", len(samp_ids))\nprint(\"overlap ids          :\", len(sub_ids & samp_ids))\nprint(\"missing from sub     :\", len(samp_ids - sub_ids))\nprint(\"extra in sub         :\", len(sub_ids - samp_ids))\n\nprint(\"example missing:\", list(samp_ids - sub_ids)[:5])\nprint(\"example extra  :\", list(sub_ids - samp_ids)[:5])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:29.112165Z","iopub.execute_input":"2026-01-18T13:35:29.112590Z","iopub.status.idle":"2026-01-18T13:35:30.214059Z","shell.execute_reply.started":"2026-01-18T13:35:29.112567Z","shell.execute_reply":"2026-01-18T13:35:30.213223Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re\nimport pandas as pd\n\nSAMPLE_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\"\nGO_RE = re.compile(r\"^GO:\\d{7}$\")\n\nsamp_raw = pd.read_csv(\n    SAMPLE_PATH,\n    sep=\"\\t\",\n    header=None,\n    engine=\"python\",\n    on_bad_lines=\"skip\"   # <-- key fix\n)\n\nsamp_raw = samp_raw.iloc[:, :3]\nsamp_raw.columns = [\"protein_id\", \"go_id\", \"score\"]\n\nsamp = samp_raw[samp_raw[\"go_id\"].astype(str).str.match(GO_RE)].copy()\n\nsamp[\"protein_id\"] = samp[\"protein_id\"].astype(str)\nsamp[\"go_id\"] = samp[\"go_id\"].astype(str)\nsamp[\"score\"] = pd.to_numeric(samp[\"score\"], errors=\"coerce\").fillna(0.0)\n\nprint(\"sample rows:\", len(samp),\n      \"unique proteins:\", samp[\"protein_id\"].nunique(),\n      \"unique go:\", samp[\"go_id\"].nunique())\nprint(samp.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:30.215164Z","iopub.execute_input":"2026-01-18T13:35:30.215642Z","iopub.status.idle":"2026-01-18T13:35:30.293106Z","shell.execute_reply.started":"2026-01-18T13:35:30.215616Z","shell.execute_reply":"2026-01-18T13:35:30.292503Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SUB_PATH = \"/kaggle/working/submission.tsv\"\nsub = pd.read_csv(SUB_PATH, sep=\"\\t\", header=None, names=[\"protein_id\",\"go_id\",\"score\"])\nsub[\"protein_id\"] = sub[\"protein_id\"].astype(str)\n\nsub_ids  = set(sub[\"protein_id\"])\nsamp_ids = set(samp[\"protein_id\"])\n\nprint(\"submission unique ids:\", len(sub_ids))\nprint(\"sample unique ids    :\", len(samp_ids))\nprint(\"overlap ids          :\", len(sub_ids & samp_ids))\nprint(\"missing from sub     :\", len(samp_ids - sub_ids))\nprint(\"extra in sub         :\", len(sub_ids - samp_ids))\nprint(\"example missing:\", list(samp_ids - sub_ids)[:5])\nprint(\"example extra  :\", list(sub_ids - samp_ids)[:5])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:30.294031Z","iopub.execute_input":"2026-01-18T13:35:30.294298Z","iopub.status.idle":"2026-01-18T13:35:31.471089Z","shell.execute_reply.started":"2026-01-18T13:35:30.294266Z","shell.execute_reply":"2026-01-18T13:35:31.470260Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SUB_PATH = \"/kaggle/working/submission.tsv\"\nsub = pd.read_csv(SUB_PATH, sep=\"\\t\", header=None, names=[\"protein_id\",\"go_id\",\"score\"])\nsub[\"protein_id\"] = sub[\"protein_id\"].astype(str)\n\ntest_ids = set(test_seq[\"protein_id\"].astype(str))\nsub_ids  = set(sub[\"protein_id\"])\n\nprint(\"missing test ids in submission:\", len(test_ids - sub_ids))\nprint(\"extra ids not in test:\", len(sub_ids - test_ids))\nprint(\"example missing:\", list(test_ids - sub_ids)[:5])\nprint(\"example extra:\", list(sub_ids - test_ids)[:5])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:31.472062Z","iopub.execute_input":"2026-01-18T13:35:31.472345Z","iopub.status.idle":"2026-01-18T13:35:32.676258Z","shell.execute_reply.started":"2026-01-18T13:35:31.472324Z","shell.execute_reply":"2026-01-18T13:35:32.675433Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import re, numpy as np\n\nGO_OK = sub[\"go_id\"].astype(str).str.match(r\"^GO:\\d{7}$\")\nprint(\"bad GO rows:\", (~GO_OK).sum())\n\nsub[\"score\"] = pd.to_numeric(sub[\"score\"], errors=\"coerce\")\nprint(\"NaN scores:\", sub[\"score\"].isna().sum())\nprint(\"score min/max:\", float(sub[\"score\"].min()), float(sub[\"score\"].max()))\nprint(\"out of [0,1]:\", ((sub[\"score\"] < 0) | (sub[\"score\"] > 1)).sum())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:32.677357Z","iopub.execute_input":"2026-01-18T13:35:32.677601Z","iopub.status.idle":"2026-01-18T13:35:33.696060Z","shell.execute_reply.started":"2026-01-18T13:35:32.677579Z","shell.execute_reply":"2026-01-18T13:35:33.695168Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dup = sub.duplicated([\"protein_id\",\"go_id\"]).sum()\nprint(\"duplicate (protein_id, go_id) rows:\", dup)\n\n# If dup > 0, fix it like this:\nsub = sub.groupby([\"protein_id\",\"go_id\"], as_index=False)[\"score\"].max()\nsub.to_csv(SUB_PATH, sep=\"\\t\", header=False, index=False)\nprint(\"rewrote without duplicates:\", SUB_PATH, \"rows:\", len(sub))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:33.697132Z","iopub.execute_input":"2026-01-18T13:35:33.697610Z","iopub.status.idle":"2026-01-18T13:35:39.579020Z","shell.execute_reply.started":"2026-01-18T13:35:33.697583Z","shell.execute_reply":"2026-01-18T13:35:39.578251Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"improving score hopefully","metadata":{}},{"cell_type":"code","source":"import re\nfrom collections import defaultdict, deque\nimport numpy as np\nimport torch\n\nASPECTS = [\"BP\", \"MF\", \"CC\"]\n\nTOPK_PER_ASP = {\"BP\": 60, \"MF\": 60, \"CC\": 60}\nMAX_ROWS_PER_PROT = 250\nBATCH_PRED = 2048 \n\nDEFAULT_FALLBACK_SCORE = 0.05\nOUT_PATH = \"/kaggle/working/submission.tsv\"\n\nOBO_PATH = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n\nID_RE  = re.compile(r\"^id:\\s+(GO:\\d{7})\")\nISA_RE = re.compile(r\"^is_a:\\s+(GO:\\d{7})\")\n\nparents = defaultdict(list)\ncur = None\nwith open(OBO_PATH, \"r\", encoding=\"utf-8\") as f:\n    for line in f:\n        line = line.strip()\n        m = ID_RE.match(line)\n        if m:\n            cur = m.group(1)\n            continue\n        if cur is not None:\n            m = ISA_RE.match(line)\n            if m:\n                parents[cur].append(m.group(1))\n\nanc_cache = {}\ndef get_ancestors(go_id: str):\n    a = anc_cache.get(go_id)\n    if a is not None:\n        return a\n    out = set()\n    dq = deque([go_id])\n    while dq:\n        x = dq.popleft()\n        for p in parents.get(x, []):\n            if p not in out:\n                out.add(p)\n                dq.append(p)\n    anc_cache[go_id] = out\n    return out\n\nprint(\"GO graph ready. #terms with parents:\", len(parents))\n\nXt = torch.tensor(X_test, dtype=torch.float32)\ntest_ids = test_pids_aligned\n\nmissing_ids = list(set(test_seq[\"protein_id\"]) - set(test_pids_aligned))\nprint(\"Missing test proteins (fallback):\", missing_ids)\n\nwritten = 0\nwith open(OUT_PATH, \"w\", encoding=\"utf-8\") as f:\n    for start in range(0, Xt.shape[0], BATCH_PRED):\n        xb = Xt[start:start+BATCH_PRED].to(device, non_blocking=True)\n        bsz = xb.shape[0]\n        batch_pids = test_ids[start:start+bsz]\n\n        per_prot = [defaultdict(float) for _ in range(bsz)]\n\n        with torch.no_grad():\n            for asp in ASPECTS:\n                model = models[asp].to(device).eval()\n                i2go = go_vocab[asp]\n                nlab = len(i2go)\n                k = min(TOPK_PER_ASP[asp], nlab)\n\n                prob = torch.sigmoid(model(xb)).float().cpu() \n                sc, idx = torch.topk(prob, k=k, dim=1)\n                sc = sc.numpy()\n                idx = idx.numpy()\n\n                for r in range(bsz):\n                    d = per_prot[r]\n                    for j in range(k):\n                        go = i2go[int(idx[r, j])]\n                        s  = float(sc[r, j])\n                        if s > d[go]:\n                            d[go] = s\n\n        for r in range(bsz):\n            d = per_prot[r]\n            items = list(d.items())\n            for go, s in items:\n                for anc in get_ancestors(go):\n                    if s > d[anc]:\n                        d[anc] = s\n\n            if len(d) > MAX_ROWS_PER_PROT:\n                top = sorted(d.items(), key=lambda x: x[1], reverse=True)[:MAX_ROWS_PER_PROT]\n                d.clear()\n                d.update(top)\n\n            pid = batch_pids[r]\n            for go, s in d.items():\n                f.write(f\"{pid}\\t{go}\\t{s:.6f}\\n\")\n            written += len(d)\n\n        if (start // BATCH_PRED) % 10 == 0:\n            print(f\"batch {start}/{Xt.shape[0]} written_rows={written}\")\n\n    for pid in missing_ids:\n        for asp in ASPECTS:\n            for go in top_fallback[asp][:200]:\n                f.write(f\"{pid}\\t{go}\\t{DEFAULT_FALLBACK_SCORE:.6f}\\n\")\n                written += 1\n\nprint(\"DONE. Wrote:\", OUT_PATH, \"| rows:\", written)\n\nbad_go = 0\nbad_score = 0\nline_count = 0\ngo_re = re.compile(r\"^GO:\\d{7}$\")\n\nwith open(OUT_PATH, \"r\", encoding=\"utf-8\") as fin:\n    for i, line in enumerate(fin):\n        line_count += 1\n        if i < 5:\n            print(\"sample:\", line.strip())\n        parts = line.rstrip(\"\\n\").split(\"\\t\")\n        if len(parts) != 3:\n            bad_score += 1\n            continue\n        pid, go, sc = parts\n        if not go_re.match(go):\n            bad_go += 1\n        try:\n            s = float(sc)\n            if not (0.0 <= s <= 1.0):\n                bad_score += 1\n        except:\n            bad_score += 1\n\nprint(\"lines:\", line_count, \"| bad_go:\", bad_go, \"| bad_score/format:\", bad_score)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:35:39.580152Z","iopub.execute_input":"2026-01-18T13:35:39.580420Z","iopub.status.idle":"2026-01-18T13:39:56.760151Z","shell.execute_reply.started":"2026-01-18T13:35:39.580388Z","shell.execute_reply":"2026-01-18T13:39:56.759280Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"have models:\", \"models\" in globals())\nprint(\"have go_vocab:\", \"go_vocab\" in globals())\nprint(\"have thr_aspect:\", \"thr_aspect\" in globals())\nprint(\"have X_test:\", \"X_test\" in globals())\nprint(\"have test_pids_aligned:\", \"test_pids_aligned\" in globals())\nprint(\"have top_fallback:\", \"top_fallback\" in globals())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-18T13:42:16.912267Z","iopub.execute_input":"2026-01-18T13:42:16.912641Z","iopub.status.idle":"2026-01-18T13:42:16.917910Z","shell.execute_reply.started":"2026-01-18T13:42:16.912615Z","shell.execute_reply":"2026-01-18T13:42:16.917085Z"}},"outputs":[],"execution_count":null}]}