{"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":14875579,"sourceType":"competition"},{"sourceId":13881193,"sourceType":"datasetVersion","datasetId":8843931},{"sourceId":13889016,"sourceType":"datasetVersion","datasetId":8848545},{"sourceId":14034995,"sourceType":"datasetVersion","datasetId":8937468}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Cài đặt môi trường và thư viện","metadata":{}},{"cell_type":"code","source":"!pip install Bio obonet networkx --quiet\n\nimport os\nimport gc\nimport random\nimport numpy as np\nimport pandas as pd\n\nfrom Bio import SeqIO\nimport obonet\nimport networkx as nx\n\nfrom tqdm.auto import tqdm\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\n\nimport torch\nimport torch.nn as nn\nfrom torch import amp\nfrom torch.utils.data import Dataset, DataLoader, random_split\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Giai đoạn 1: Tối ưu hóa dự đoán bằng GOA UniProt","metadata":{}},{"cell_type":"code","source":"FILE_MLP = \"/kaggle/input/submission-mlp/submission.tsv\"\nGOA_FILE = \"/kaggle/input/protein-go-annotations/goa_uniprot_all.csv\"\n\ngraph = obonet.read_obo(\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n)\nreversed_graph = graph.reverse() # Mặc định OBO là Con->Cha. Cần đảo thành Cha->Con để tìm hậu duệ\n\ndescendants_map = {}\nfor node in tqdm(graph.nodes, desc=\"Mapping Descendants\"):\n    descendants_map[node] = nx.descendants(reversed_graph, node) # Tìm con cháu\n\nnegative_keys = set()\npositive_keys = set()\npositive_rows = []\n\nfor chunk in tqdm(pd.read_csv(GOA_FILE, usecols=['protein_id', 'go_term', 'qualifier'], chunksize=1000000), desc=\"Scanning GOA\"):\n    # Lọc các dòng có NOT\n    neg_chunk = chunk[chunk['qualifier'].astype(str).str.contains('NOT', na=False)]\n\n    # zip tránh ghép sai cặp\n    for protein_id, go_term in zip(neg_chunk['protein_id'], neg_chunk['go_term']):\n        negative_keys.add(f\"{protein_id}_{go_term}\")\n\n        # Nếu Protein ko có chức năng cha -> ko có chức năng con\n        if go_term in descendants_map:\n            for desc in descendants_map[go_term]:\n                negative_keys.add(f\"{protein_id}_{desc}\") \n\n    # Lọc các dòng ko có NOT\n    pos_chunk = chunk[~chunk['qualifier'].astype(str).str.contains('NOT', na=False)]\n    pos_chunk = pos_chunk.drop_duplicates(subset=['protein_id', 'go_term'])\n\n    for protein_id, go_term in zip(pos_chunk['protein_id'], pos_chunk['go_term']):\n        key = f\"{protein_id}_{go_term}\"\n        if key not in positive_keys:\n            positive_rows.append(f\"{protein_id}\\t{go_term}\\t1.000\")\n            positive_keys.add(key)\n\n# Ghi file submission1\nwith open(\"submission1.tsv\", \"w\") as f_out:\n    # Ghi dữ liệu chuẩn theo UniProt\n    if positive_rows:\n        f_out.write(\"\\n\".join(positive_rows) + \"\\n\")\n    # Chống tràn ram\n    del positive_rows\n    gc.collect()\n\n    # Ghi dữ liệu từ submission-mlp\n    with open(FILE_MLP, \"r\") as f_in:\n        for line in tqdm(f_in, desc=\"Reading MLP\"):\n            parts = line.strip().split('\\t')\n            if len(parts) < 3:\n                continue\n\n            protein_id, go_term, score = parts[0], parts[1], parts[2]\n            key = f\"{protein_id}_{go_term}\"\n\n            # Nếu protein có NOT hoặc protein đã ghi ở trên rồi thì ko ghi nữa\n            if key not in negative_keys and key not in positive_keys:\n                f_out.write(line)\n\nprint(\"Done\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Giai đoạn 2: Xây dựng Mô hình Residual MLP","metadata":{}},{"cell_type":"markdown","source":"### Chuẩn bị dữ liệu huấn luyện và test","metadata":{}},{"cell_type":"code","source":"INPUT_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\nEMBED_DIR = \"/kaggle/input/cafa6-protein-embeddings-esm2\"\n\n# ids[i] là protein tương ứng với emb[i]\nemb = np.load(f\"{EMBED_DIR}/protein_embeddings.npy\")\nids = pd.read_csv(f\"{EMBED_DIR}/protein_ids.csv\")[\"protein_id\"].tolist()\n\n# Lấy id của protein trong tập train\ntrain_df = pd.read_csv(f\"{INPUT_DIR}/Train/train_terms.tsv\", sep=\"\\t\")\ntrain_ids = train_df[\"EntryID\"].unique().tolist()\n\n# Lấy id của protein trong tập test\ntest_ids = []\nfor protein in SeqIO.parse(f\"{INPUT_DIR}/Test/testsuperset.fasta\", \"fasta\"):\n    test_ids.append(protein.id.split()[0])\n\n# Lấy index của các protein trong emb\npid_to_index = {pid: i for i, pid in enumerate(ids)}\ntrain_index = [pid_to_index[pid] for pid in train_ids if pid in pid_to_index]\ntest_index = [pid_to_index[pid] for pid in test_ids if pid in pid_to_index]\n\n# Lấy emb protein của train và test \nX_train = emb[train_index]\nX_test = emb[test_index]\n\n# Gom các nhãn GO của từng protein lại thành list\nlabels = train_df.groupby(\"EntryID\")[\"term\"].apply(list)\n\n# Lấy id của protein vừa có embedding vừa có nhãn GO\nprotein_ids = [pid for pid in train_ids if pid in labels.index]\n\n# Multi-hot Encoding\nmlb = MultiLabelBinarizer()\nY_train = mlb.fit_transform(labels[protein_ids])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Xây dựng Dataset và mô hình Residual MLP","metadata":{}},{"cell_type":"code","source":"class ProteinDataset(Dataset):\n    def __init__(self, X, Y):\n        self.X = torch.tensor(X, dtype=torch.float16 if torch.cuda.is_available() else torch.float32)\n        self.Y = torch.tensor(Y, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.X)\n\n    def __getitem__(self, i):\n        return self.X[i], self.Y[i]\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_features, hidden_features, dropout=0.3):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.BatchNorm1d(in_features),\n            nn.ReLU(),\n            nn.Linear(in_features, hidden_features),\n            nn.BatchNorm1d(hidden_features),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_features, in_features) \n        )\n        \n    def forward(self, x):\n        return x + self.block(x) \n\nclass ResNetMLP(nn.Module):\n    def __init__(self, input_dim=1280, output_dim=None):\n        super().__init__()\n        self.stem = nn.Sequential(\n            nn.Linear(input_dim, 2048),\n            nn.BatchNorm1d(2048),\n            nn.ReLU(),\n            nn.Dropout(0.2)\n        )\n        self.layer1 = ResidualBlock(2048, 1024, dropout=0.3)\n        self.layer2 = ResidualBlock(2048, 1024, dropout=0.3)\n        self.layer3 = ResidualBlock(2048, 1024, dropout=0.4)\n        self.head = nn.Linear(2048, output_dim)\n\n    def forward(self, x):\n        x = self.stem(x)\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n        return self.head(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Huấn luyện mô hình ","metadata":{}},{"cell_type":"code","source":"# Khóa random seed để kết quả huấn luyện có thể tái lập\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\nseed_everything(42)\n\nbatch_size = 256\nepochs = 18\nlr = 2e-3\npos_weight_val = 2.0\n\ndataset = ProteinDataset(X_train, Y_train)\nval_size = int(0.1 * len(dataset))\ntrain_size = len(dataset) - val_size\ntrain_ds, val_ds = random_split(dataset, [train_size, val_size])\n\n# Chia batch, shuffle=True: mỗi epoch sẽ xáo trộn dữ liệu\ntrain_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2)\nval_loader = DataLoader(val_ds, batch_size=batch_size, num_workers=2)\n\n# Khởi tạo model\nmodel = ResNetMLP(input_dim=X_train.shape[1], output_dim=Y_train.shape[1]).to(device)\n\npos_weight = torch.ones([Y_train.shape[1]]).to(device) * pos_weight_val\ncriterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\noptimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-3) \nscheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(\n    optimizer, T_0=5, T_mult=2, eta_min=1e-5\n)\nscaler = amp.GradScaler(device=\"cuda\")\nbest_val_loss = float('inf') # khởi tạo với vô cùng\n\nprint(\"Start training...\")\nfor epoch in range(epochs):\n    # Huấn luyện\n    model.train()\n    train_loss = 0\n    for Xb, Yb in train_loader:\n        Xb, Yb = Xb.to(device), Yb.to(device)\n        optimizer.zero_grad()\n        with amp.autocast(\"cuda\"):\n            logits = model(Xb)\n            loss = criterion(logits, Yb)\n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        train_loss += loss.item() * Xb.size(0)\n    train_loss /= len(train_loader.dataset)\n\n    # Đánh giá trên tập val\n    model.eval()\n    val_loss = 0\n    with torch.no_grad(): # Không tính gradient\n        for Xb, Yb in val_loader:\n            Xb, Yb = Xb.to(device), Yb.to(device)\n            with amp.autocast(\"cuda\"):\n                logits = model(Xb)\n                loss = criterion(logits, Yb)\n            val_loss += loss.item() * Xb.size(0)\n    val_loss /= len(val_loader.dataset)\n\n    # Cập nhật learning rate\n    scheduler.step(epoch + val_loss/100)\n    lr_current = optimizer.param_groups[0]['lr']\n\n    print(f\"Epoch {epoch+1}/{epochs} | Train Loss: {train_loss:.5f} | Val Loss: {val_loss:.5f} | LR: {lr_current:.1e}\")\n\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        torch.save(model.state_dict(), \"protein_resnet_best.pt\")\n        print(f\"--> Save Best Model (Val Loss: {val_loss:.5f})\")\n    else:\n        print(\"--> No improvement.\")\n\nprint(f\"\\n Training completed. Best Val Loss: {best_val_loss:.5f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Dự đoán trên tập test","metadata":{}},{"cell_type":"code","source":"torch.cuda.empty_cache()\n\n# Khởi tạo mô hình\nmodel = ResNetMLP(input_dim=X_train.shape[1], output_dim=Y_train.shape[1]).to(device)\n# Load trọng số từ best model bên trên\nmodel.load_state_dict(torch.load(\"protein_resnet_best.pt\"))\nmodel.eval() \nmodel = model.half()\n\ntest_id_to_index = {pid: i for i, pid in enumerate(test_ids)}\nall_go_terms = mlb.classes_\nbatch_size = 32\nmin_prob = 0.01\ntop_k = 25\n\nsubmission_rows = []\nfor i in tqdm(range(0, len(test_ids), batch_size), desc=\"Predicting...\"):\n    batch_ids = test_ids[i:i+batch_size]\n    batch_idx = [test_id_to_index[pid] for pid in batch_ids]\n    batch = X_test[batch_idx].astype(np.float16)\n    batch = torch.tensor(batch, dtype=torch.float16).to(device)\n\n    # Tính xác suất dự đoán cho mỗi nhãn GO\n    with torch.no_grad():\n        logits = model(batch)\n        probs = torch.sigmoid(logits).cpu().numpy()\n\n    for j, pid in enumerate(batch_ids):\n        p = probs[j]\n        idx = np.where(p >= min_prob)[0]\n        if top_k is not None and len(idx) < top_k:\n            top_extra = p.argsort()[-top_k:]\n            idx = np.unique(np.concatenate([idx, top_extra]))\n        idx = idx[np.argsort(-p[idx])]\n        if top_k is not None:\n             idx = idx[:top_k]\n\n        for k in idx:\n            term = all_go_terms[k]\n            score = float(p[k])\n            submission_rows.append([pid, term, score])\n\n    del batch, logits, probs\n    torch.cuda.empty_cache()\n\n# Ghi file submission2\nsubmission_df = pd.DataFrame(submission_rows, columns=[\"EntryID\", \"GO\", \"Score\"])\nsubmission_df.to_csv(\"submission2.tsv\", sep=\"\\t\", index=False, header=False)\nprint(\"Done\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Giai đoạn 3: Ensemble","metadata":{}},{"cell_type":"code","source":"WEIGHT_A = 0.7\nWEIGHT_B = 0.3\nDIVERSITY_BONUS = 0.05\nTHRESHOLD = 0.01\nTOP_K = 30\nFILE_A = \"/kaggle/working/submission1.tsv\"\nFILE_B = \"/kaggle/working/submission2.tsv\"\n\n# Lấy các go terms hợp lệ\ntrain_df = pd.read_csv(f\"{INPUT_DIR}/Train/train_terms.tsv\", sep=\"\\t\", usecols=[\"term\"])\nvalid_terms = set(train_df[\"term\"].unique())\n\npreds = {}\nprint(\"Load file A...\")\nwith open(FILE_A, 'r') as f:\n    for line in tqdm(f):\n        parts = line.strip().split('\\t')\n        if len(parts) < 3: continue       \n        if parts[1] not in valid_terms: continue         \n        key = f\"{parts[0]}\\t{parts[1]}\"\n        preds[key] = {'a': float(parts[2]), 'b': 0.0}\n        \nprint(\"Loading file B...\")\nwith open(FILE_B, 'r') as f:\n    for line in tqdm(f):\n        parts = line.strip().split('\\t')\n        if len(parts) < 3: continue  \n        if parts[1] not in valid_terms: continue    \n        key = f\"{parts[0]}\\t{parts[1]}\"   \n        if key in preds: preds[key]['b'] = float(parts[2])\n        else: preds[key] = {'a': 0.0, 'b': float(parts[2])}\n\nfinal_data = []\nfor key, scores in tqdm(preds.items(), desc=\"Computing...\"):\n    score_a = scores['a']\n    score_b = scores['b']\n    final_score = (score_a * WEIGHT_A) + (score_b * WEIGHT_B)    \n    if score_a > 0.01 and score_b > 0.01:\n        final_score += DIVERSITY_BONUS # Thêm điểm bonus nếu cả hai đều dự đoán xác suất > 0.01\n    final_score = min(final_score, 1.0)   \n    if final_score >= THRESHOLD:\n        pid, term = key.split('\\t')\n        final_data.append([pid, term, final_score])\ndel preds\n\n# Ghi file submission3\ndf = pd.DataFrame(final_data, columns=[\"Id\", \"Term\", \"Score\"])\ndf = df.sort_values(by=[\"Id\", \"Score\"], ascending=[True, False])\ndf_clean = df.groupby(\"Id\").head(TOP_K).copy()\ndf_clean[\"Score\"] = df_clean[\"Score\"].round(3)\ndf_clean.to_csv(\"submission3.tsv\", sep='\\t', index=False, header=False)\n\nprint(\"Done\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Link tải file submission của 3 giai đoạn","metadata":{}},{"cell_type":"code","source":"from IPython.display import FileLink\n\nprint(\"Giai đoạn 1 (0.32):\")\ndisplay(FileLink(\"submission1.tsv\"))\nprint(\"Giai đoạn 2 (0.295):\")\ndisplay(FileLink(\"submission2.tsv\"))\nprint(\"Giai đoạn 3 (0.342):\")\ndisplay(FileLink(\"submission3.tsv\"))","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}