{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session\n\n# Use the kagglehub client library to attach Kaggle resources like competitions, datasets, and models to your session\n# Learn more about kagglehub: https://github.com/Kaggle/kagglehub/blob/main/README.md\n\nimport kagglehub\n# kagglehub.dataset_download('<owner>/<dataset-slug>')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-06-28T18:42:38.439387Z","iopub.execute_input":"2026-06-28T18:42:38.439696Z","iopub.status.idle":"2026-06-28T18:42:40.540419Z","shell.execute_reply.started":"2026-06-28T18:42:38.439656Z","shell.execute_reply":"2026-06-28T18:42:40.539489Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, glob, subprocess, shutil\nfrom collections import defaultdict\nimport numpy as np\nimport pandas as pd\nimport torch, torch.nn as nn\nfrom sklearn.model_selection import train_test_split\nimport gc\n\nDATA = \"/kaggle/input/competitions/cafa-6-protein-function-prediction\"\nWORK = \"/kaggle/working\"\nTRAIN_FASTA = f\"{DATA}/Train/train_sequences.fasta\"\nTRAIN_TERMS = f\"{DATA}/Train/train_terms.tsv\"\nOBO         = f\"{DATA}/Train/go-basic.obo\"\nTEST_FASTA  = f\"{DATA}/Test/testsuperset.fasta\"\nTAXON = f\"{DATA}/Test/testsuperset-taxon-list.tsv\"\nIA_FILE     = f\"{DATA}/IA.txt\"\nASPECTS = [\"BPO\", \"CCO\", \"MFO\"]\nTOP_TERMS = {\"BPO\": 1500, \"CCO\": 800, \"MFO\": 1000}\nMAX_PER_PROTEIN = 1500\nSEED = 42\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"device:\", DEVICE, \"| data exists:\", os.path.exists(TRAIN_FASTA))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def parse_obo(path):\n    ns, parents = {}, defaultdict(set)\n    cur = cur_ns = None; cur_par = set(); obsolete = False\n    def flush():\n        if cur and not obsolete:\n            if cur_ns: ns[cur] = cur_ns\n            parents[cur] |= cur_par\n    with open(path, encoding=\"utf-8\") as fh:\n        for line in fh:\n            line = line.rstrip(\"\\n\")\n            if line == \"[Term]\" or (line.startswith(\"[\") and line.endswith(\"]\")):\n                flush(); cur = cur_ns = None; cur_par = set(); obsolete = False\n            elif line.startswith(\"id: GO:\"):       cur = line[4:].strip()\n            elif line.startswith(\"namespace: \"):    cur_ns = line[11:].strip()\n            elif line.startswith(\"is_obsolete: true\"): obsolete = True\n            elif line.startswith(\"is_a: GO:\"):       cur_par.add(line.split()[1].strip())\n            elif line.startswith(\"relationship: part_of GO:\"): cur_par.add(line.split()[2].strip())\n        flush()\n    return ns, dict(parents)\n\n\nclass GODag:\n    def __init__(self, path):\n        self.ns, self.parents = parse_obo(path); self._c = {}\n    def ancestors(self, t):\n        if t in self._c: return self._c[t]\n        seen, stack = set(), [t]\n        while stack:\n            x = stack.pop()\n            if x in seen: continue\n            seen.add(x); stack.extend(self.parents.get(x, ()))\n        self._c[t] = seen; return seen\n    def propagate_terms(self, terms):\n        out = set()\n        for t in terms: out |= self.ancestors(t)\n        return out\n    def ancestor_pairs(self, vocab):\n        idx = {t: i for i, t in enumerate(vocab)}; pairs = []\n        for t, ci in idx.items():\n            for a in self.ancestors(t):\n                ai = idx.get(a)\n                if ai is not None: pairs.append((ci, ai))\n        return np.asarray(pairs, dtype=np.int64)\n    @staticmethod\n    def propagate_scores(scores, pairs):\n        out = scores.copy()\n        np.maximum.at(out, (slice(None), pairs[:, 1]), scores[:, pairs[:, 0]])\n        return out\n\n\ndef weighted_fmax(y_true, y_score, ia=None, thr=None):\n    y_true = np.asarray(y_true, float); y_score = np.asarray(y_score, float)\n    V = y_true.shape[1]\n    ia = np.ones(V) if ia is None else np.asarray(ia, float)\n    thr = np.round(np.arange(0.01, 1.0, 0.01), 2) if thr is None else thr\n    w_true = y_true * ia; t_w = w_true.sum(1); n_truth = int((t_w > 0).sum())\n    if n_truth == 0: return 0.0\n    best = 0.0\n    for t in thr:\n        pred = y_score >= t\n        pp = (pred * ia).sum(1); tp = (pred * w_true).sum(1)\n        hp = pp > 0; m = int(hp.sum())\n        if m == 0: continue\n        pr = (tp[hp] / pp[hp]).sum() / m\n        rc = (tp[t_w > 0] / t_w[t_w > 0]).sum() / n_truth\n        f = 0.0 if pr + rc == 0 else 2 * pr * rc / (pr + rc)\n        best = max(best, f)\n    return best","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def read_fasta(path):\n    seqs, pid, buf = {}, None, []\n    with open(path, encoding=\"utf-8\") as fh:\n        for line in fh:\n            line = line.rstrip()\n            if line.startswith(\">\"):\n                if pid is not None: seqs[pid] = \"\".join(buf)\n                pid = line[1:].split()[0]; buf = []\n            else: buf.append(line)\n        if pid is not None: seqs[pid] = \"\".join(buf)\n    return seqs\n\n\ndef load_terms():\n    df = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\n    df.columns = [c.strip() for c in df.columns]\n    ren = {}\n    for c in df.columns:\n        cl = c.lower()\n        if cl.startswith(\"entry\"): ren[c] = \"EntryID\"\n        elif cl == \"term\" or cl.startswith(\"go\"): ren[c] = \"term\"\n        elif cl.startswith(\"aspect\"): ren[c] = \"aspect\"\n    return df.rename(columns=ren)[[\"EntryID\", \"term\", \"aspect\"]]\n\n\ndef build_vocab(terms, dag, aspect):\n    sub = terms[terms.aspect == aspect]; counts = defaultdict(int)\n    for _, grp in sub.groupby(\"EntryID\")[\"term\"]:\n        for t in dag.propagate_terms(grp): counts[t] += 1\n    return sorted(counts, key=counts.get, reverse=True)[:TOP_TERMS[aspect]]\n\n\ndef label_matrix(terms, dag, aspect, vocab, proteins):\n    vi = {t: i for i, t in enumerate(vocab)}; pi = {p: i for i, p in enumerate(proteins)}\n    Y = np.zeros((len(proteins), len(vocab)), np.float32)\n    sub = terms[terms.aspect == aspect]\n    for pid, grp in sub.groupby(\"EntryID\")[\"term\"]:\n        i = pi.get(pid)\n        if i is None: continue\n        for t in dag.propagate_terms(grp):\n            j = vi.get(t)\n            if j is not None: Y[i, j] = 1.0\n    return Y\n\n\ndef load_ia(vocab):\n    w = np.ones(len(vocab))\n    if not os.path.exists(IA_FILE): return w\n    ia = {}\n    with open(IA_FILE) as fh:\n        for line in fh:\n            p = line.split()\n            if len(p) >= 2:\n                try: ia[p[0]] = float(p[1])\n                except ValueError: pass\n    for i, t in enumerate(vocab):\n        if t in ia: w[i] = ia[t]\n    return w\n\n\ndef write_submission(path, proteins, blocks, min_score=1e-3):\n    per = [dict() for _ in proteins]\n    for vocab, sc in blocks:\n        for j, term in enumerate(vocab):\n            col = sc[:, j]\n            for i in np.where(col >= min_score)[0]:\n                s = float(col[i])\n                if s > per[i].get(term, 0.0): per[i][term] = s\n    with open(path, \"w\") as fh:\n        for i, pid in enumerate(proteins):\n            for term, s in sorted(per[i].items(), key=lambda kv: -kv[1])[:MAX_PER_PROTEIN]:\n                fh.write(f\"{pid}\\t{term}\\t{round(s, 3)}\\n\")\n    print(\"wrote\", path)\n\n\nprint(\"loading DAG + data ...\")\ndag = GODag(OBO)\nterms = load_terms()\ntrain_seqs = read_fasta(TRAIN_FASTA)\ntest_seqs = read_fasta(TEST_FASTA)\ntrain_ids = sorted(train_seqs); test_ids = sorted(test_seqs)\nprint(f\"train={len(train_ids)} test={len(test_ids)} terms={len(terms)}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd         \n# esm-2 plm\ndef load_aligned_by_index(emb_path, id_path, target_ids):\n    X = np.load(emb_path).astype(np.float32)\n    original_ids = np.load(id_path, allow_pickle=True)\n\n    original_ids = list(original_ids)\n\n    id_to_index = {orig_id: index for index, orig_id in enumerate(original_ids)}\n\n    ordered_indices = [id_to_index[t_id] for t_id in target_ids]\n\n    X_aligned = X[ordered_indices]\n\n    # L2 正規化\n    return X_aligned / (np.linalg.norm(X_aligned, axis=1, keepdims=True) + 1e-8)\n\n\nXtr = load_aligned_by_index(f\"/kaggle/input/datasets/anoreo/cafa6-t5-embeddings/train_embeddings_esm2.npy\", \"/kaggle/input/datasets/anoreo/cafa6-t5-embeddings/train_ids_esm2.npy\",train_ids)\nXte = load_aligned_by_index(f\"/kaggle/input/datasets/anoreo/cafa6-t5-embeddings/test_embeddings_esm2.npy\",  \"/kaggle/input/datasets/anoreo/cafa6-t5-embeddings/test_ids_esm2.npy\",test_ids)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-06-29T15:32:43.515606Z","iopub.execute_input":"2026-06-29T15:32:43.515989Z","iopub.status.idle":"2026-06-29T15:35:50.649744Z","shell.execute_reply.started":"2026-06-29T15:32:43.515946Z","shell.execute_reply":"2026-06-29T15:35:50.648886Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nclass MLP(nn.Module):\n    def __init__(self, d_in, d_out, h=1024, p=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(d_in, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(p),\n            nn.Linear(h, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(p),\n            nn.Linear(h, d_out))\n    def forward(self, x): return self.net(x)\n\ndef train_mlp(Xt, Yt, Xpred, epochs=20, bs=256, lr=1e-3):\n    m = MLP(Xt.shape[1], Yt.shape[1]).to(DEVICE)\n    opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=1e-5)\n    lf = nn.BCEWithLogitsLoss()\n    Xt_t = torch.tensor(Xt, device=DEVICE); Yt_t = torch.tensor(Yt, device=DEVICE)\n    for _ in range(epochs):\n        m.train(); perm = torch.randperm(len(Xt_t), device=DEVICE)\n        for i in range(0, len(Xt_t), bs):\n            idx = perm[i:i + bs]; opt.zero_grad()\n            lf(m(Xt_t[idx]), Yt_t[idx]).backward(); opt.step()\n    m.eval()\n    with torch.no_grad():\n        return torch.sigmoid(m(torch.tensor(Xpred, device=DEVICE))).cpu().numpy()\n\nesm2_mlp_blocks = []\nfor aspect in ASPECTS:\n    vocab = build_vocab(terms, dag, aspect)\n    ia = load_ia(vocab)\n    # 取得階層關係\n    pairs = dag.ancestor_pairs(vocab)\n\n    # 如果 pairs 是空的，這會將它的 shape 從 (0,) 變成安全的 (0, 2)\n    pairs = np.array(pairs).reshape(-1, 2)\n    \n    Y = label_matrix(terms, dag, aspect, vocab, train_ids)\n    tr, va = train_test_split(np.arange(len(train_ids)), test_size=0.1, random_state=SEED)\n    \n    pva = GODag.propagate_scores(train_mlp(Xtr[tr], Y[tr], Xtr[va]), pairs)\n    print(f\"[{aspect}] Fmax(mlp)={weighted_fmax(Y[va], pva, ia):.4f}\")\n    \n    pte = GODag.propagate_scores(train_mlp(Xtr, Y, Xte), pairs)\n    esm2_mlp_blocks.append((vocab, pte.astype(np.float32)))\n\nwrite_submission(f\"{WORK}/submission_esm2_mlp.tsv\", test_ids, esm2_mlp_blocks)\n\ntry:\n    del Xtr, Xte, Y\nexcept NameError:\n    pass \n\ngc.collect()\n\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# t5 plm\n\ndef load_aligned_by_index(emb_path, id_path, target_ids):\n    X = np.load(emb_path).astype(np.float32)\n    original_ids = np.load(id_path, allow_pickle=True)\n\n    original_ids = list(original_ids)\n\n    id_to_index = {orig_id: index for index, orig_id in enumerate(original_ids)}\n\n    ordered_indices = [id_to_index[t_id] for t_id in target_ids]\n\n    X_aligned = X[ordered_indices]\n\n    # L2 正規化\n    return X_aligned / (np.linalg.norm(X_aligned, axis=1, keepdims=True) + 1e-8)\n\n\nXtr = load_aligned_by_index(f\"/kaggle/input/datasets/kaitonmh/cafa-6-embeddings/protT5_xl/train_sequences_emb.npy\", \"/kaggle/input/datasets/kaitonmh/cafa-6-embeddings/protT5_xl/train_sequences_ids.npy\",train_ids)\nXte = load_aligned_by_index(f\"/kaggle/input/datasets/kaitonmh/cafa-6-embeddings/protT5_xl/testsuperset_emb.npy\",  \"/kaggle/input/datasets/kaitonmh/cafa-6-embeddings/protT5_xl/testsuperset_ids.npy\",test_ids)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nclass MLP(nn.Module):\n    def __init__(self, d_in, d_out, h=1024, p=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(d_in, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(p),\n            nn.Linear(h, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(p),\n            nn.Linear(h, d_out))\n    def forward(self, x): return self.net(x)\n\ndef train_mlp(Xt, Yt, Xpred, epochs=20, bs=256, lr=1e-3):\n    m = MLP(Xt.shape[1], Yt.shape[1]).to(DEVICE)\n    opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=1e-5)\n    lf = nn.BCEWithLogitsLoss()\n    Xt_t = torch.tensor(Xt, device=DEVICE); Yt_t = torch.tensor(Yt, device=DEVICE)\n    for _ in range(epochs):\n        m.train(); perm = torch.randperm(len(Xt_t), device=DEVICE)\n        for i in range(0, len(Xt_t), bs):\n            idx = perm[i:i + bs]; opt.zero_grad()\n            lf(m(Xt_t[idx]), Yt_t[idx]).backward(); opt.step()\n    m.eval()\n    with torch.no_grad():\n        return torch.sigmoid(m(torch.tensor(Xpred, device=DEVICE))).cpu().numpy()\n\nprotbert_mlp_blocks = []\nfor aspect in ASPECTS:\n    vocab = build_vocab(terms, dag, aspect)\n    ia = load_ia(vocab)\n\n    pairs = dag.ancestor_pairs(vocab)\n    pairs = np.array(pairs).reshape(-1, 2)\n    \n    Y = label_matrix(terms, dag, aspect, vocab, train_ids)\n    tr, va = train_test_split(np.arange(len(train_ids)), test_size=0.1, random_state=SEED)\n    \n    pva = GODag.propagate_scores(train_mlp(Xtr[tr], Y[tr], Xtr[va]), pairs)\n    print(f\"[{aspect}] Fmax(mlp)={weighted_fmax(Y[va], pva, ia):.4f}\")\n    \n    pte = GODag.propagate_scores(train_mlp(Xtr, Y, Xte), pairs)\n    protbert_mlp_blocks.append((vocab, pte.astype(np.float32)))\n    \nwrite_submission(f\"{WORK}/submission_protbert_mlp.tsv\", test_ids, protbert_mlp_blocks)\n\ntry:\n    del Xtr, Xte, Y\nexcept NameError:\n    pass \n\ngc.collect()\n\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!apt-get update && apt-get install -y diamond-aligner","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %% [CELL 6] (optional) DIAMOND homology label transfer -------------------\ndef diamond_homology():\n    db = f\"{WORK}/traindb\"\n    subprocess.run([\"diamond\", \"makedb\", \"--in\", TRAIN_FASTA, \"-d\", db], check=True)\n    subprocess.run([\"diamond\", \"blastp\", \"-q\", TEST_FASTA, \"-d\", db,\n                    \"-o\", f\"{WORK}/hits.tsv\", \"--outfmt\", \"6\", \"qseqid\", \"sseqid\",\n                    \"bitscore\", \"--max-target-seqs\", \"250\", \"--evalue\", \"0.001\"], check=True)\n    hits = pd.read_csv(f\"{WORK}/hits.tsv\", sep=\"\\t\", header=None, names=[\"q\", \"s\", \"bit\"])\n    tidx = {p: i for i, p in enumerate(test_ids)}; out = []\n    for aspect in ASPECTS:\n        vocab = build_vocab(terms, dag, aspect); vi = {t: j for j, t in enumerate(vocab)}\n        sub = terms[terms.aspect == aspect]; pt = {}\n        for pid, grp in sub.groupby(\"EntryID\")[\"term\"]:\n            pt[pid] = [vi[t] for t in dag.propagate_terms(grp) if t in vi]\n        sc = np.zeros((len(test_ids), len(vocab)), np.float32); norm = np.zeros(len(test_ids), np.float32)\n        for q, s, bit in hits.itertuples(index=False):\n            i = tidx.get(q); cols = pt.get(s)\n            if i is None or not cols: continue\n            sc[i, cols] += bit; norm[i] += bit\n        norm[norm == 0] = 1.0; sc /= norm[:, None]\n        out.append((vocab, sc))\n    return out\n\nhom_blocks = diamond_homology()   ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# %% [CELL 7] ensemble + final submission ----------------------------------\n# taxon prior replaces the plain global naive as the cheap \"prior\" source.\nW = {\"mlp\": 0.75, \"homology\": 0.25, \"taxon\": 0.10}\nfinal = []\nfor k, aspect in enumerate(ASPECTS):\n    vocab = build_vocab(terms, dag, aspect); pairs = dag.ancestor_pairs(vocab)\n    pairs = dag.ancestor_pairs(vocab)\n    pairs = np.array(pairs).reshape(-1, 2)\n    acc = W[\"mlp\"] * (protbert_mlp_blocks[k][1] + esm2_mlp_blocks[k][1]) /2\n    wsum = W[\"mlp\"]\n    # taxon prior (CELL 3b)\n    try:\n        acc = acc + W[\"taxon\"] * tax_blocks[k][1]; wsum += W[\"taxon\"]\n    except NameError:\n        pass\n    # homology (CELL 6)\n    try:\n        acc = acc + W[\"homology\"] * hom_blocks[k][1]; wsum += W[\"homology\"]\n    except NameError:\n        pass\n\n    acc = GODag.propagate_scores(acc / wsum, pairs)\n    final.append((vocab, acc))\n\nwrite_submission(f\"{WORK}/submission.tsv\", test_ids, final, min_score=0.01)\n\n# 驗證:必須看到「沒有外層雙引號」的三欄\nprint(\"\\n 最終檔案格式預覽：\")\nwith open(f\"{WORK}/submission.tsv\") as f:\n    for _ in range(5):\n        print(repr(f.readline().rstrip(\"\\n\")))","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}