{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":14058983,"sourceType":"datasetVersion","datasetId":8948309}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install fair-esm pandas scikit-learn tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:33:10.457744Z","iopub.execute_input":"2025-12-10T12:33:10.457988Z","iopub.status.idle":"2025-12-10T12:33:16.758799Z","shell.execute_reply.started":"2025-12-10T12:33:10.457959Z","shell.execute_reply":"2025-12-10T12:33:16.757902Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport pandas as pd\n\nDATA_DIR = \"/kaggle/input/data-test\"   # đổi đúng tên dataset của bạn\nWORK_DIR = \"/kaggle/working\"\n\nTRAIN_DIR = os.path.join(DATA_DIR, \"Train\")\n\nTERMS_PATH = os.path.join(TRAIN_DIR, \"train_terms.tsv\")\nFASTA_PATH = os.path.join(TRAIN_DIR, \"train_sequences.fasta\")\n\nprint(\"DATA_DIR:\", DATA_DIR)\nprint(\"TRAIN_DIR:\", TRAIN_DIR)\nprint(\"Có gì trong Train:\", os.listdir(TRAIN_DIR))\n\ndef parse_fasta_ids(fasta_path):\n    seqs = {}\n    with open(fasta_path, \"r\") as f:\n        seq_id = None\n        parts = []\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if seq_id is not None:\n                    seqs[seq_id] = \"\".join(parts)\n                header = line[1:]\n                if \"|\" in header:\n                    fields = header.split(\"|\")\n                    if len(fields) >= 2:\n                        seq_id = fields[1]\n                    else:\n                        seq_id = header.split()[0]\n                else:\n                    seq_id = header.split()[0]\n                parts = []\n            else:\n                parts.append(line)\n        if seq_id is not None:\n            seqs[seq_id] = \"\".join(parts)\n    return seqs\n\nseq_dict = parse_fasta_ids(FASTA_PATH)\nprint(\"Số sequence đọc được:\", len(seq_dict))","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:39:39.936699Z","iopub.execute_input":"2025-12-10T12:39:39.937460Z","iopub.status.idle":"2025-12-10T12:39:41.972234Z","shell.execute_reply.started":"2025-12-10T12:39:39.937421Z","shell.execute_reply":"2025-12-10T12:39:41.971429Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from collections import defaultdict\n\nOBO_PATH = os.path.join(TRAIN_DIR, \"go-basic.obo\")\nprint(\"OBO_PATH:\", OBO_PATH, \"exists:\", os.path.exists(OBO_PATH))\n\nparents = defaultdict(set)  # child -> set(parent)\nnamespace = {}              # term -> 'F' / 'P' / 'C'\n\nwith open(OBO_PATH, \"r\") as f:\n    current_id = None\n    in_term = False\n    for line in f:\n        line = line.strip()\n        if line == \"[Term]\":\n            in_term = True\n            current_id = None\n            continue\n        if line == \"\" and in_term:\n            in_term = False\n            current_id = None\n            continue\n        if not in_term:\n            continue\n\n        if line.startswith(\"id: GO:\"):\n            current_id = line.split(\"id: \")[1]\n        elif line.startswith(\"is_a: GO:\") and current_id is not None:\n            pid = line.split(\"is_a: \")[1].split()[0]\n            parents[current_id].add(pid)\n        elif line.startswith(\"relationship: part_of GO:\") and current_id is not None:\n            pid = line.split(\"relationship: part_of \")[1].split()[0]\n            parents[current_id].add(pid)\n        elif line.startswith(\"namespace:\") and current_id is not None:\n            ns = line.split(\"namespace: \")[1]\n            if ns == \"molecular_function\":\n                namespace[current_id] = \"F\"\n            elif ns == \"biological_process\":\n                namespace[current_id] = \"P\"\n            elif ns == \"cellular_component\":\n                namespace[current_id] = \"C\"\n\nprint(\"Số term có parent:\", len(parents))\nprint(\"Số term có namespace:\", len(namespace))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:40:09.612864Z","iopub.execute_input":"2025-12-10T12:40:09.613873Z","iopub.status.idle":"2025-12-10T12:40:10.503822Z","shell.execute_reply.started":"2025-12-10T12:40:09.613844Z","shell.execute_reply":"2025-12-10T12:40:10.503090Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from functools import lru_cache\n\n@lru_cache(maxsize=None)\ndef get_ancestors(term):\n    \"\"\"Trả về {term} ∪ tất cả tổ tiên của nó.\"\"\"\n    result = {term}\n    for p in parents.get(term, []):\n        result |= get_ancestors(p)\n    return result\n\nTERMS_PATH = os.path.join(TRAIN_DIR, \"train_terms.tsv\")\nterms = pd.read_csv(TERMS_PATH, sep=\"\\t\")  # EntryID, term, aspect\nprint(\"Số dòng train_terms:\", len(terms))\n\ndef build_train_for_aspect(aspect, out_csv):\n    sub = terms[terms[\"aspect\"] == aspect].copy()\n    print(\"\\nAspect\", aspect, \"số dòng:\", len(sub))\n\n    grouped = sub.groupby(\"EntryID\")[\"term\"].apply(set).reset_index(name=\"terms_raw\")\n\n    rows = []\n    for _, row in grouped.iterrows():\n        pid = row[\"EntryID\"]\n        if pid not in seq_dict:\n            continue\n        seq = seq_dict[pid]\n\n        label_set = set()\n        for t in row[\"terms_raw\"]:\n            anc = get_ancestors(t)\n            for a in anc:\n                if namespace.get(a) == aspect:  # giữ đúng ontology\n                    label_set.add(a)\n\n        if not label_set:\n            continue\n\n        labels = \";\".join(sorted(label_set))\n        rows.append((pid, seq, labels))\n\n    df_out = pd.DataFrame(rows, columns=[\"protein_id\", \"sequence\", \"labels\"])\n    df_out.to_csv(out_csv, index=False)\n    print(\"Saved\", aspect, \"->\", out_csv, \"rows:\", len(df_out))\n\nbuild_train_for_aspect(\"F\", os.path.join(WORK_DIR, \"train_MF.csv\"))\nbuild_train_for_aspect(\"P\", os.path.join(WORK_DIR, \"train_BP.csv\"))\nbuild_train_for_aspect(\"C\", os.path.join(WORK_DIR, \"train_CC.csv\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:40:32.497995Z","iopub.execute_input":"2025-12-10T12:40:32.498527Z","iopub.status.idle":"2025-12-10T12:40:48.728134Z","shell.execute_reply.started":"2025-12-10T12:40:32.498504Z","shell.execute_reply":"2025-12-10T12:40:48.727398Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom torch.utils.data import Dataset\n\nclass ProteinFunctionDataset(Dataset):\n    def __init__(self, csv_path, go_terms_list=None):\n        self.df = pd.read_csv(csv_path)\n        self.df['label_list'] = self.df['labels'].fillna(\"\").apply(\n            lambda x: [lab.strip() for lab in x.split(';') if lab.strip() != \"\"]\n        )\n        \n        if go_terms_list is None:\n            all_terms = set()\n            for labs in self.df['label_list']:\n                all_terms.update(labs)\n            self.go_terms = sorted(list(all_terms))\n        else:\n            self.go_terms = go_terms_list\n        \n        self.term2idx = {t: i for i, t in enumerate(self.go_terms)}\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        seq = row['sequence']\n        labels = row['label_list']\n        \n        y = torch.zeros(len(self.go_terms), dtype=torch.float32)\n        for lab in labels:\n            if lab in self.term2idx:\n                y[self.term2idx[lab]] = 1.0\n        \n        return row['protein_id'], seq, y\n\nds_MF = ProteinFunctionDataset(os.path.join(WORK_DIR, \"train_MF.csv\"))\nds_BP = ProteinFunctionDataset(os.path.join(WORK_DIR, \"train_BP.csv\"))\nds_CC = ProteinFunctionDataset(os.path.join(WORK_DIR, \"train_CC.csv\"))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:46:56.156686Z","iopub.execute_input":"2025-12-10T12:46:56.156973Z","iopub.status.idle":"2025-12-10T12:46:58.951949Z","shell.execute_reply.started":"2025-12-10T12:46:56.156952Z","shell.execute_reply":"2025-12-10T12:46:58.951317Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import esm\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\nclass ESMEmbedder:\n    def __init__(self, model_name=\"esm2_t12_35M_UR50D\", device=None, max_len=1022):\n        self.device = device or (\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.model, self.alphabet = esm.pretrained.__dict__[model_name]()\n        self.model.eval()\n        self.model.to(self.device)\n        self.batch_converter = self.alphabet.get_batch_converter()\n        self.max_len = max_len   # giới hạn chiều dài\n\n    @torch.no_grad()\n    def encode_batch(self, ids, seqs):\n        # Cắt chuỗi nếu quá dài\n        clipped_seqs = [s[:self.max_len] for s in seqs]\n\n        batch_data = list(zip(ids, clipped_seqs))\n        labels, strs, toks = self.batch_converter(batch_data)\n        toks = toks.to(self.device)\n\n        out = self.model(toks, repr_layers=[self.model.num_layers])\n        token_reps = out[\"representations\"][self.model.num_layers]\n\n        emb_dict = {}\n        for i, pid in enumerate(ids):\n            seq_len = (toks[i] != self.alphabet.padding_idx).sum().item()\n            rep = token_reps[i, 1:seq_len-1].mean(0)\n            emb_dict[pid] = rep.cpu()\n        return emb_dict\n\ndef precompute_embeddings(dataset, output_path, batch_size=16, model_name=\"esm2_t12_35M_UR50D\"):\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    embedder = ESMEmbedder(model_name=model_name, device=device)\n    \n    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=lambda b: b)\n    all_embs = {}\n    for batch in tqdm(loader, desc=\"Computing ESM embeddings\"):\n        ids = [x[0] for x in batch]\n        seqs = [x[1] for x in batch]\n        emb_dict = embedder.encode_batch(ids, seqs)\n        all_embs.update(emb_dict)\n    \n    torch.save(all_embs, output_path)\n    print(\"Saved:\", output_path)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:47:24.126870Z","iopub.execute_input":"2025-12-10T12:47:24.127594Z","iopub.status.idle":"2025-12-10T12:47:24.145927Z","shell.execute_reply.started":"2025-12-10T12:47:24.127555Z","shell.execute_reply":"2025-12-10T12:47:24.145436Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport os\n\n# Gộp 3 file train_MF/BP/CC thành 1 file all (mỗi protein chỉ giữ 1 lần)\nmf_path = os.path.join(WORK_DIR, \"train_MF.csv\")\nbp_path = os.path.join(WORK_DIR, \"train_BP.csv\")\ncc_path = os.path.join(WORK_DIR, \"train_CC.csv\")\n\ndf_mf = pd.read_csv(mf_path)\ndf_bp = pd.read_csv(bp_path)\ndf_cc = pd.read_csv(cc_path)\n\ndf_all = pd.concat([df_mf, df_bp, df_cc], ignore_index=True)\n\n# Nếu cùng protein xuất hiện ở nhiều file thì giữ 1 dòng (sequence giống nhau)\ndf_all = df_all.drop_duplicates(subset=\"protein_id\")\n\nALL_TRAIN_PATH = os.path.join(WORK_DIR, \"train_all.csv\")\ndf_all.to_csv(ALL_TRAIN_PATH, index=False)\nprint(\"Saved train_all.csv:\", ALL_TRAIN_PATH, \"rows:\", len(df_all))\n\n# Tính embedding từ file all này\ntrain_dataset = ProteinFunctionDataset(ALL_TRAIN_PATH)\nEMB_PATH = os.path.join(WORK_DIR, \"train_esm_emb.pt\")\n\nprecompute_embeddings(train_dataset, EMB_PATH, batch_size=1, model_name=\"esm2_t12_35M_UR50D\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T12:50:04.085885Z","iopub.execute_input":"2025-12-10T12:50:04.086601Z","iopub.status.idle":"2025-12-10T13:29:35.055280Z","shell.execute_reply.started":"2025-12-10T12:50:04.086556Z","shell.execute_reply":"2025-12-10T13:29:35.054564Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader, random_split\nimport torch.nn as nn\nfrom torch.optim import AdamW\nfrom sklearn.metrics import f1_score\n\nclass EmbeddingDataset(Dataset):\n    def __init__(self, base_dataset, emb_path):\n        self.base_dataset = base_dataset\n        self.emb_dict = torch.load(emb_path)\n    \n    def __len__(self):\n        return len(self.base_dataset)\n    \n    def __getitem__(self, idx):\n        pid, seq, y = self.base_dataset[idx]\n        emb = self.emb_dict[pid]\n        return emb, y\n\nclass MLPClassifier(nn.Module):\n    def __init__(self, emb_dim, num_labels, hidden_dim=1024, dropout=0.4):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(emb_dim, hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, num_labels),\n        )\n    def forward(self, x):\n        return self.net(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T13:29:43.519236Z","iopub.execute_input":"2025-12-10T13:29:43.519655Z","iopub.status.idle":"2025-12-10T13:29:44.556837Z","shell.execute_reply.started":"2025-12-10T13:29:43.519631Z","shell.execute_reply":"2025-12-10T13:29:44.556268Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader, random_split\nimport torch.nn as nn\nfrom torch.optim import AdamW\nfrom sklearn.metrics import f1_score\nfrom tqdm import tqdm\nimport os\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nprint(\"Device:\", device)\n\nEMB_PATH = os.path.join(WORK_DIR, \"train_esm_emb.pt\")\n\ndef train_one_ontology(csv_path, emb_path, ckpt_path, num_epochs=10):\n    print(\"\\n=== Training ontology from:\", csv_path, \"===\")\n    base_dataset = ProteinFunctionDataset(csv_path)\n    go_terms = base_dataset.go_terms\n    emb_dataset = EmbeddingDataset(base_dataset, emb_path)\n\n    val_ratio = 0.1\n    n_total = len(emb_dataset)\n    n_val = int(n_total * val_ratio)\n    n_train = n_total - n_val\n    train_ds, val_ds = random_split(emb_dataset, [n_train, n_val])\n\n    batch_size = 64\n    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)\n\n    sample_emb, _ = emb_dataset[0]\n    emb_dim = sample_emb.shape[0]\n    num_labels = len(go_terms)\n\n    model = MLPClassifier(emb_dim, num_labels).to(device)\n    optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2)\n    criterion = nn.BCEWithLogitsLoss()\n\n    best_f1 = -1.0\n    best_state = None\n\n    for epoch in range(1, num_epochs+1):\n        # ---- train ----\n        model.train()\n        tot_loss = 0.0\n        for embs, labels in tqdm(train_loader, desc=f\"Epoch {epoch} [Train]\"):\n            embs = embs.to(device)\n            labels = labels.to(device)\n            optimizer.zero_grad()\n            logits = model(embs)\n            loss = criterion(logits, labels)\n            loss.backward()\n            optimizer.step()\n            tot_loss += loss.item() * embs.size(0)\n        train_loss = tot_loss / n_train\n\n        # ---- val ----\n        model.eval()\n        tot_loss = 0.0\n        all_t, all_p = [], []\n        with torch.no_grad():\n            for embs, labels in tqdm(val_loader, desc=f\"Epoch {epoch} [Val]\"):\n                embs = embs.to(device)\n                labels = labels.to(device)\n                logits = model(embs)\n                loss = criterion(logits, labels)\n                tot_loss += loss.item() * embs.size(0)\n                probs = torch.sigmoid(logits)\n                preds = (probs > 0.5).float()\n                all_t.append(labels.cpu())\n                all_p.append(preds.cpu())\n        val_loss = tot_loss / n_val\n        all_t = torch.cat(all_t).numpy()\n        all_p = torch.cat(all_p).numpy()\n        f1_micro = f1_score(all_t, all_p, average=\"micro\", zero_division=0)\n        print(f\"Epoch {epoch}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}, F1_micro={f1_micro:.4f}\")\n\n        if f1_micro > best_f1:\n            best_f1 = f1_micro\n            best_state = model.state_dict()\n            print(f\"--> New best F1_micro={best_f1:.4f} at epoch {epoch}\")\n\n    torch.save({\n        \"model_state\": best_state,\n        \"go_terms\": go_terms,\n        \"emb_dim\": emb_dim,\n    }, ckpt_path)\n    print(\"Saved best model to\", ckpt_path, \"with F1_micro =\", best_f1)\n\n\n# Gọi train cho 3 ontology\ntrain_one_ontology(os.path.join(WORK_DIR, \"train_MF.csv\"), EMB_PATH,\n                   os.path.join(WORK_DIR, \"mlp_MF.pt\"))\n\ntrain_one_ontology(os.path.join(WORK_DIR, \"train_BP.csv\"), EMB_PATH,\n                   os.path.join(WORK_DIR, \"mlp_BP.pt\"))\n\ntrain_one_ontology(os.path.join(WORK_DIR, \"train_CC.csv\"), EMB_PATH,\n                   os.path.join(WORK_DIR, \"mlp_CC.pt\"))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T13:33:30.832935Z","iopub.execute_input":"2025-12-10T13:33:30.833405Z","iopub.status.idle":"2025-12-10T13:48:41.415402Z","shell.execute_reply.started":"2025-12-10T13:33:30.833377Z","shell.execute_reply":"2025-12-10T13:48:41.414615Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom tqdm import tqdm\n\nclass MultiOntologyPredictor:\n    def __init__(\n        self,\n        ckpt_MF,\n        ckpt_BP,\n        ckpt_CC,\n        csv_MF,\n        csv_BP,\n        csv_CC,\n        emb_path,\n        model_name=\"esm2_t12_35M_UR50D\",\n        device=None,\n        knn_k=20,      # số hàng xóm gần nhất\n        alpha=0.6,     # trọng số MLP (0.6) vs kNN (0.4)\n    ):\n        self.device = device or (\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.knn_k = knn_k\n        self.alpha = alpha\n\n        # ----- 1) Load 3 checkpoint MLP -----\n        self.ck_MF = torch.load(ckpt_MF, map_location=\"cpu\")\n        self.ck_BP = torch.load(ckpt_BP, map_location=\"cpu\")\n        self.ck_CC = torch.load(ckpt_CC, map_location=\"cpu\")\n\n        self.go_terms_MF = self.ck_MF[\"go_terms\"]\n        self.go_terms_BP = self.ck_BP[\"go_terms\"]\n        self.go_terms_CC = self.ck_CC[\"go_terms\"]\n\n        emb_dim = self.ck_MF[\"emb_dim\"]  # 3 model dùng chung emb_dim\n\n        self.model_MF = MLPClassifier(emb_dim, len(self.go_terms_MF))\n        self.model_BP = MLPClassifier(emb_dim, len(self.go_terms_BP))\n        self.model_CC = MLPClassifier(emb_dim, len(self.go_terms_CC))\n\n        self.model_MF.load_state_dict(self.ck_MF[\"model_state\"])\n        self.model_BP.load_state_dict(self.ck_BP[\"model_state\"])\n        self.model_CC.load_state_dict(self.ck_CC[\"model_state\"])\n\n        for m in [self.model_MF, self.model_BP, self.model_CC]:\n            m.to(self.device)\n            m.eval()\n\n        # ----- 2) Load toàn bộ embedding train -----\n        print(\"Loading train embeddings from:\", emb_path)\n        emb_dict_all = torch.load(emb_path, map_location=\"cpu\")\n\n        # ----- 3) Xây dữ liệu kNN cho từng ontology -----\n        def build_knn_data(csv_path, go_terms, name):\n            print(f\"Build kNN data for {name} from {csv_path}\")\n            ds = ProteinFunctionDataset(csv_path, go_terms_list=go_terms)\n            embs = []\n            labels = []\n            for i in tqdm(range(len(ds)), desc=f\"kNN {name}\"):\n                pid, seq, y = ds[i]\n                if pid in emb_dict_all:\n                    embs.append(emb_dict_all[pid])\n                    labels.append(y)\n            embs = torch.stack(embs)           # (N, emb_dim) trên CPU\n            labels = torch.stack(labels).float()  # (N, num_labels) CPU\n            # chuẩn hoá để dùng cosine\n            embs = embs / embs.norm(dim=1, keepdim=True).clamp(min=1e-8)\n            # embedding dùng cho tính similarity để trên GPU, labels để CPU\n            return embs.to(self.device), labels\n\n        self.knn_embs_MF, self.knn_labels_MF = build_knn_data(csv_MF, self.go_terms_MF, \"MF\")\n        self.knn_embs_BP, self.knn_labels_BP = build_knn_data(csv_BP, self.go_terms_BP, \"BP\")\n        self.knn_embs_CC, self.knn_labels_CC = build_knn_data(csv_CC, self.go_terms_CC, \"CC\")\n\n        # ----- 4) Embedder ESM dùng chung -----\n        self.embedder = ESMEmbedder(model_name=model_name, device=self.device)\n\n    def _knn_scores(self, emb_norm, knn_embs, knn_labels):\n        \"\"\"\n        emb_norm: (1, D) trên device (đã chuẩn hoá)\n        knn_embs: (N, D) trên device (đã chuẩn hoá)\n        knn_labels: (N, T) trên CPU (multi-hot)\n        return: (T,) vector score kNN trên CPU\n        \"\"\"\n        sims = torch.mm(emb_norm, knn_embs.T).squeeze(0)   # (N,)\n        k = min(self.knn_k, sims.size(0))\n        vals, idx = torch.topk(sims, k=k)\n\n        # bỏ similarity âm, tránh làm loãng\n        vals = torch.clamp(vals, min=0)\n        if vals.sum().item() <= 0:\n            return torch.zeros(knn_labels.size(1))\n\n        weights = (vals / vals.sum()).cpu()          # (k,)\n        chosen_labels = knn_labels[idx.cpu()]        # (k, T)\n        scores = torch.matmul(weights.unsqueeze(0), chosen_labels).squeeze(0)  # (T,)\n        return scores\n\n    @torch.no_grad()\n    def predict(self, protein_id, sequence, top_k=50):\n        # ----- 1) Tính embedding ESM -----\n        emb_dict = self.embedder.encode_batch([protein_id], [sequence])\n        emb = emb_dict[protein_id].unsqueeze(0).to(self.device)    # (1, D)\n        emb_norm = emb / emb.norm(dim=1, keepdim=True).clamp(min=1e-8)\n\n        # ----- 2) MLP + kNN cho từng ontology -----\n        def run_one(model, terms, knn_embs, knn_labels, name):\n            logits = model(emb)\n            p_mlp = torch.sigmoid(logits).squeeze(0).cpu()          # (T,)\n            p_knn = self._knn_scores(emb_norm, knn_embs, knn_labels)  # (T,) CPU\n            p_final = self.alpha * p_mlp + (1 - self.alpha) * p_knn  # (T,)\n            return [\n                {\"go_term\": t, \"prob\": float(p)}\n                for t, p in zip(terms, p_final.numpy())\n            ]\n\n        preds_MF = run_one(self.model_MF, self.go_terms_MF,\n                           self.knn_embs_MF, self.knn_labels_MF, \"MF\")\n        preds_BP = run_one(self.model_BP, self.go_terms_BP,\n                           self.knn_embs_BP, self.knn_labels_BP, \"BP\")\n        preds_CC = run_one(self.model_CC, self.go_terms_CC,\n                           self.knn_embs_CC, self.knn_labels_CC, \"CC\")\n\n        # ----- 3) Gộp & chọn top_k -----\n        all_preds = preds_MF + preds_BP + preds_CC\n        all_preds.sort(key=lambda x: x[\"prob\"], reverse=True)\n        if top_k is not None:\n            all_preds = all_preds[:top_k]\n        return all_preds\n\n\n# Khởi tạo predictor 3 ontology + kNN\nCK_MF = os.path.join(WORK_DIR, \"mlp_MF.pt\")\nCK_BP = os.path.join(WORK_DIR, \"mlp_BP.pt\")\nCK_CC = os.path.join(WORK_DIR, \"mlp_CC.pt\")\n\nCSV_MF = os.path.join(WORK_DIR, \"train_MF.csv\")\nCSV_BP = os.path.join(WORK_DIR, \"train_BP.csv\")\nCSV_CC = os.path.join(WORK_DIR, \"train_CC.csv\")\n\nEMB_PATH = os.path.join(WORK_DIR, \"train_esm_emb.pt\")\n\nmulti_pred = MultiOntologyPredictor(\n    CK_MF, CK_BP, CK_CC,\n    CSV_MF, CSV_BP, CSV_CC,\n    EMB_PATH,\n    knn_k=20,\n    alpha=0.6,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T17:16:21.851729Z","iopub.execute_input":"2025-12-10T17:16:21.852059Z","iopub.status.idle":"2025-12-10T17:17:06.641679Z","shell.execute_reply.started":"2025-12-10T17:16:21.852036Z","shell.execute_reply":"2025-12-10T17:17:06.640547Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom tqdm import tqdm\n\nTEST_FASTA = os.path.join(DATA_DIR, \"Test\", \"testsuperset.fasta\")\nSUB_PATH = os.path.join(WORK_DIR, \"submission.tsv\")\n\ndef parse_test_fasta(fasta_path):\n    \"\"\"Yield (protein_id, sequence) từ testsuperset.fasta\"\"\"\n    with open(fasta_path, \"r\") as f:\n        seq_id = None\n        parts = []\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                # trả về seq cũ\n                if seq_id is not None:\n                    yield seq_id, \"\".join(parts)\n                header = line[1:]\n                if \"|\" in header:\n                    fields = header.split(\"|\")\n                    if len(fields) >= 2:\n                        seq_id = fields[1]\n                    else:\n                        seq_id = header.split()[0]\n                else:\n                    seq_id = header.split()[0]\n                parts = []\n            else:\n                parts.append(line)\n        # seq cuối\n        if seq_id is not None:\n            yield seq_id, \"\".join(parts)\n\ndef count_fasta_sequences(fasta_path):\n    \"\"\"Đếm tổng số protein trong file FASTA (số header '>')\"\"\"\n    n = 0\n    with open(fasta_path, \"r\") as f:\n        for line in f:\n            if line.startswith(\">\"):\n                n += 1\n    return n\n\ndef make_submission_3ont(test_fasta_path, output_path,\n                         predictor, top_k=50, min_prob=0.01):\n    \"\"\"Tạo file submission.tsv: protein_id \\t go_term \\t probability\"\"\"\n    total_proteins = count_fasta_sequences(test_fasta_path)\n    print(\"Tổng số protein test:\", total_proteins)\n\n    with open(output_path, \"w\") as out:\n        for pid, seq in tqdm(\n            parse_test_fasta(test_fasta_path),\n            total=total_proteins,\n            desc=\"Processing proteins\"\n        ):\n            preds = predictor.predict(pid, seq, top_k=top_k)\n            for p in preds:\n                if p[\"prob\"] >= min_prob:\n                    out.write(f\"{pid}\\t{p['go_term']}\\t{p['prob']:.6f}\\n\")\n\n    print(\"Đã lưu submission ở:\", output_path)\n\nmake_submission_3ont(TEST_FASTA, SUB_PATH, multi_pred, top_k=50, min_prob=0.01)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-10T17:18:53.119190Z","iopub.execute_input":"2025-12-10T17:18:53.119522Z","iopub.status.idle":"2025-12-10T20:07:47.980948Z","shell.execute_reply.started":"2025-12-10T17:18:53.119498Z","shell.execute_reply":"2025-12-10T20:07:47.980226Z"}},"outputs":[],"execution_count":null}]}