{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":13665986,"sourceType":"datasetVersion","datasetId":8688768}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q transformers obonet biopython torchmetrics","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import KFold\nfrom tqdm.auto import tqdm\nimport os\nimport gc\nimport obonet\nimport networkx as nx\n\n# ==========================================\n# 1. CẤU HÌNH PATH (THEO HÌNH ẢNH CỦA BẠN)\n# ==========================================\nclass Config:\n    # Root Dir (Tên folder dataset chính trên Kaggle)\n    # Lưu ý: Tên folder có thể khác nhau tùy lúc Add data, hãy check kỹ bên phải notebook\n    RAW_DATA_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n    EMBED_DIR = '/kaggle/input/cafa-6-t5-embeddings'\n    \n    # Cấu trúc folder con (Dựa trên hình ảnh 2)\n    TRAIN_DIR = os.path.join(RAW_DATA_DIR, 'Train')\n    TEST_DIR = os.path.join(RAW_DATA_DIR, 'Test')\n    \n    # File Paths\n    TRAIN_TERMS = os.path.join(TRAIN_DIR, 'train_terms.tsv')\n    TRAIN_TAXONOMY = os.path.join(TRAIN_DIR, 'train_taxonomy.tsv')\n    OBO_FILE = os.path.join(TRAIN_DIR, 'go-basic.obo')\n    \n    # Test Taxonomy (Quan trọng để tăng điểm Test set)\n    TEST_TAXONOMY = os.path.join(TEST_DIR, 'testsuperset-taxon-list.tsv')\n    \n    # Embeddings Paths (Dựa trên hình ảnh 1)\n    TRAIN_EMBEDS = os.path.join(EMBED_DIR, 'train_embeds.npy')\n    TRAIN_IDS = os.path.join(EMBED_DIR, 'train_ids.npy')\n    TEST_EMBEDS = os.path.join(EMBED_DIR, 'test_embeds.npy')\n    TEST_IDS = os.path.join(EMBED_DIR, 'test_ids.npy')\n    \n    # Hyperparameters\n    NUM_CLASSES = 1500      # Số lượng GO Terms (Labels)\n    BATCH_SIZE = 256        # Batch lớn cho GPU chạy nhanh\n    EPOCHS = 12             # Số epoch mỗi fold\n    LR = 1e-3\n    N_FOLDS = 5             # Ensemble 5 mô hình\n    TOP_TAXONS_COUNT = 50   # Top 50 loài phổ biến nhất để one-hot\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ncfg = Config()\nprint(f\"Running on: {cfg.DEVICE}\")\nprint(f\"Train Taxonomy Path: {cfg.TRAIN_TAXONOMY}\")\n\n# ==========================================\n# 2. XỬ LÝ DỮ LIỆU & TAXONOMY\n# ==========================================\n# ==========================================\n# CẬP NHẬT: HÀM XỬ LÝ ID ĐỂ SỬA LỖI N_SAMPLES=0\n# ==========================================\n\ndef clean_ids(id_list):\n    \"\"\"\n    Chuyển đổi ID từ dạng 'sp|P12345|XYZ' sang 'P12345' nếu cần.\n    \"\"\"\n    cleaned = []\n    for x in id_list:\n        # Chuyển bytes sang string nếu cần\n        if isinstance(x, bytes):\n            x = x.decode('utf-8')\n        x = str(x).strip()\n        \n        # Xử lý format sp|Accession|Name\n        if '|' in x:\n            # Lấy phần ở giữa (Accession ID)\n            parts = x.split('|')\n            if len(parts) >= 2:\n                cleaned.append(parts[1]) # Lấy P12345\n            else:\n                cleaned.append(x)\n        else:\n            cleaned.append(x)\n    return np.array(cleaned)\n\ndef load_and_align_data_fixed():\n    print(\"--- 1. Loading Embeddings & Fixing IDs ---\")\n    X_train = np.load(cfg.TRAIN_EMBEDS)\n    train_ids_raw = np.load(cfg.TRAIN_IDS)\n    \n    X_test = np.load(cfg.TEST_EMBEDS)\n    test_ids_raw = np.load(cfg.TEST_IDS)\n    \n    # --- BƯỚC SỬA LỖI QUAN TRỌNG: CLEAN IDS ---\n    train_ids = clean_ids(train_ids_raw)\n    test_ids = clean_ids(test_ids_raw)\n    \n    # DEBUG: In ra để kiểm tra xem ID đã khớp chưa\n    print(f\"Sample Raw Train ID: {train_ids_raw[0]}\")\n    print(f\"Sample Cleaned Train ID: {train_ids[0]}\")\n    print(f\"Train shape: {X_train.shape}\")\n    \n    # --- XỬ LÝ TAXONOMY ---\n    print(\"--- 2. Processing Taxonomy Features ---\")\n    \n    # Đọc file Taxonomy\n    train_tax_df = pd.read_csv(cfg.TRAIN_TAXONOMY, sep='\\t', dtype=str)\n    # Map: ProteinID -> TaxonID\n    # Lưu ý: Cột đầu tiên của taxonomy cũng phải khớp với cleaned ID\n    train_tax_map = dict(zip(train_tax_df.iloc[:, 0], train_tax_df.iloc[:, 1]))\n    \n    test_tax_df = pd.read_csv(cfg.TEST_TAXONOMY, sep='\\t', dtype=str)\n    test_tax_map = dict(zip(test_tax_df.iloc[:, 0], test_tax_df.iloc[:, 1]))\n    \n    top_taxons = train_tax_df.iloc[:, 1].value_counts().head(cfg.TOP_TAXONS_COUNT).index.tolist()\n    tax2idx = {t: i for i, t in enumerate(top_taxons)}\n    \n    def create_tax_features(ids, tax_map):\n        feats = np.zeros((len(ids), cfg.TOP_TAXONS_COUNT), dtype=np.float32)\n        match_count = 0\n        for i, pid in enumerate(ids):\n            if pid in tax_map:\n                match_count += 1\n                tid = tax_map[pid]\n                if tid in tax2idx:\n                    feats[i, tax2idx[tid]] = 1.0\n        print(f\"   > Found taxonomy for {match_count}/{len(ids)} proteins\")\n        return feats\n\n    print(\"Generating Taxon Vectors...\")\n    train_tax_feats = create_tax_features(train_ids, train_tax_map)\n    test_tax_feats = create_tax_features(test_ids, test_tax_map)\n    \n    X_train_final = np.concatenate([X_train, train_tax_feats], axis=1)\n    X_test_final = np.concatenate([X_test, test_tax_feats], axis=1)\n    \n    return X_train_final, train_ids, X_test_final, test_ids\n\n\n# (Sau đó chạy tiếp phần Inference như cũ)\n# ==========================================\n# 3. MODEL: RESNET (Deep & Residual)\n# ==========================================\nclass ResidualBlock(nn.Module):\n    def __init__(self, dim, dropout=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.BatchNorm1d(dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(dim, dim),\n            nn.BatchNorm1d(dim)\n        )\n        self.relu = nn.ReLU()\n    def forward(self, x):\n        return self.relu(x + self.net(x)) # Skip Connection\n\nclass AdvancedModel(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super().__init__()\n        # Input Projection: Đưa vector (1074) lên không gian lớn hơn\n        self.entry = nn.Sequential(\n            nn.Linear(input_dim, 1024),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Dropout(0.2)\n        )\n        # Deep Residual Layers\n        self.blocks = nn.Sequential(\n            ResidualBlock(1024),\n            ResidualBlock(1024),\n            ResidualBlock(1024)\n        )\n        self.head = nn.Linear(1024, num_classes)\n        \n    def forward(self, x):\n        x = self.entry(x)\n        x = self.blocks(x)\n        return self.head(x)\n\n# ==========================================\n# 4. TRAINING UTILS\n# ==========================================\nclass ProteinDataset(Dataset):\n    def __init__(self, X, Y=None):\n        self.X = torch.from_numpy(X).float()\n        self.Y = torch.from_numpy(Y).float() if Y is not None else None\n    def __len__(self): return len(self.X)\n    def __getitem__(self, i): \n        return (self.X[i], self.Y[i]) if self.Y is not None else self.X[i]\n\ndef train_one_fold(fold, model, train_loader, val_loader):\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = optim.AdamW(model.parameters(), lr=cfg.LR, weight_decay=1e-4)\n    # Scheduler: Giảm LR nếu loss đi ngang\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=1)\n    \n    best_loss = float('inf')\n    early_stop_count = 0\n    patience = 3 # Dừng sớm nếu 3 epoch không cải thiện\n    \n    for epoch in range(cfg.EPOCHS):\n        model.train()\n        t_loss = 0\n        for x, y in train_loader:\n            x, y = x.to(cfg.DEVICE), y.to(cfg.DEVICE)\n            optimizer.zero_grad()\n            logits = model(x)\n            loss = criterion(logits, y)\n            loss.backward()\n            optimizer.step()\n            t_loss += loss.item()\n            \n        # Validation\n        model.eval()\n        v_loss = 0\n        with torch.no_grad():\n            for x, y in val_loader:\n                x, y = x.to(cfg.DEVICE), y.to(cfg.DEVICE)\n                v_loss += criterion(model(x), y).item()\n        \n        avg_v_loss = v_loss / len(val_loader)\n        scheduler.step(avg_v_loss)\n        \n        # Save Best Model\n        if avg_v_loss < best_loss:\n            best_loss = avg_v_loss\n            torch.save(model.state_dict(), f'model_fold_{fold}.pth')\n            early_stop_count = 0\n        else:\n            early_stop_count += 1\n            \n        if early_stop_count >= patience:\n            print(f\"  Fold {fold}: Early stopping at epoch {epoch}\")\n            break\n    \n    return best_loss\n\n# ==========================================\n# 5. MAIN PIPELINE (ENSEMBLE 5-FOLD)\n# ==========================================\n\n# --- CHẠY LẠI PHẦN LOAD DATA VỚI HÀM MỚI ---\nX_train_full, train_ids_full, X_test_full, test_ids_full = load_and_align_data_fixed()\n\nprint(\"--- 3. Processing Labels & Checking Mismatch ---\")\ndf = pd.read_csv(cfg.TRAIN_TERMS, sep='\\t', names=['EntryID', 'term', 'aspect'])\nprint(f\"Sample ID in Train Terms: {df['EntryID'].iloc[0]}\") # Debug xem ID trong file nhãn trông thế nào\n\n# Lấy Top Labels\ntop_terms = df['term'].value_counts().head(cfg.NUM_CLASSES).index.tolist()\nterm2idx = {t: i for i, t in enumerate(top_terms)}\n\n# Map ID -> Set of Terms\nid_to_terms = df[df['term'].isin(top_terms)].groupby('EntryID')['term'].apply(set).to_dict()\n\n# Tạo matrix Y\nY_full = np.zeros((len(train_ids_full), cfg.NUM_CLASSES), dtype=np.float32)\nvalid_mask = []\n\nfor i, pid in enumerate(train_ids_full):\n    if pid in id_to_terms:\n        valid_mask.append(True)\n        for term in id_to_terms[pid]:\n            Y_full[i, term2idx[term]] = 1.0\n    else:\n        valid_mask.append(False)\n\nvalid_mask = np.array(valid_mask)\nprint(f\"DEBUG: Total proteins: {len(train_ids_full)}\")\nprint(f\"DEBUG: Matched proteins with labels: {sum(valid_mask)}\")\n\nif sum(valid_mask) == 0:\n    raise ValueError(\"CRITICAL ERROR: No IDs matched! Check the 'Sample Cleaned Train ID' vs 'Sample ID in Train Terms' printed above.\")\n\n# Chỉ lấy data khớp\nX_train_clean = X_train_full[valid_mask]\nY_clean = Y_full[valid_mask]\n\n# --- TIẾP TỤC CHẠY PHẦN TRAINING LOOP CŨ TỪ ĐÂY ---\n# Copy lại đoạn \"kf = KFold...\" từ code trước và chạy tiếp\nkf = KFold(n_splits=cfg.N_FOLDS, shuffle=True, random_state=42)\ninput_dim = X_train_clean.shape[1]\n\nprint(f\"\\n--- Starting Ensemble Training ({cfg.N_FOLDS} Folds) ---\")\nfor fold, (train_idx, val_idx) in enumerate(kf.split(X_train_clean)):\n    # ... (Code training giữ nguyên như cũ) ...\n    print(f\">> Fold {fold+1}/{cfg.N_FOLDS}\")\n    \n    train_ds = ProteinDataset(X_train_clean[train_idx], Y_clean[train_idx])\n    val_ds = ProteinDataset(X_train_clean[val_idx], Y_clean[val_idx])\n    \n    train_loader = DataLoader(train_ds, batch_size=cfg.BATCH_SIZE, shuffle=True)\n    val_loader = DataLoader(val_ds, batch_size=cfg.BATCH_SIZE)\n    \n    model = AdvancedModel(input_dim, cfg.NUM_CLASSES).to(cfg.DEVICE)\n    train_one_fold(fold, model, train_loader, val_loader)\n    \n    del model, train_loader, val_loader, train_ds, val_ds\n    torch.cuda.empty_cache()\n    gc.collect()\n\n\n# C. INFERENCE & AVERAGING\nprint(\"\\n--- Running Inference on Test Set ---\")\ntest_loader = DataLoader(ProteinDataset(X_test_full), batch_size=cfg.BATCH_SIZE)\nfinal_preds = np.zeros((len(test_ids_full), cfg.NUM_CLASSES), dtype=np.float32)\n\nfor fold in range(cfg.N_FOLDS):\n    model = AdvancedModel(input_dim, cfg.NUM_CLASSES).to(cfg.DEVICE)\n    model.load_state_dict(torch.load(f'model_fold_{fold}.pth'))\n    model.eval()\n    \n    fold_preds = []\n    with torch.no_grad():\n        for x in tqdm(test_loader, desc=f\"Predicting Fold {fold+1}\"):\n            x = x.to(cfg.DEVICE)\n            logits = model(x)\n            fold_preds.append(torch.sigmoid(logits).cpu().numpy())\n            \n    final_preds += np.vstack(fold_preds)\n\nfinal_preds /= cfg.N_FOLDS # Average kết quả 5 fold\n\n# D. GO PROPAGATION (HACK ĐIỂM SỐ)\nprint(\"--- Applying GO Structure Propagation ---\")\n# Load GO Graph\ngo_graph = obonet.read_obo(cfg.OBO_FILE)\nterm_to_idx = {t: i for i, t in enumerate(top_terms)}\n\n# Tìm cha của từng term\nparent_map = {}\nvalid_set = set(top_terms)\nfor term in top_terms:\n    if term in go_graph:\n        parents = list(go_graph.successors(term))\n        for p in parents:\n            if p in valid_set:\n                if p not in parent_map: parent_map[p] = []\n                parent_map[p].append(term)\n\n# Lan truyền: Score Cha = Max(Score Cha, Max(Score Con))\n# Chạy 2 pass để đảm bảo thấm lên trên\nfor _ in range(2):\n    for parent, children in parent_map.items():\n        if children:\n            p_idx = term_to_idx[parent]\n            c_indices = [term_to_idx[c] for c in children]\n            max_child_score = final_preds[:, c_indices].max(axis=1)\n            final_preds[:, p_idx] = np.maximum(final_preds[:, p_idx], max_child_score)\n\n# E. SUBMISSION FILE\nprint(\"--- Creating Submission File ---\")\nsubmission_rows = []\nthreshold = 0.005 # Ngưỡng thấp để giữ lại nhiều kết quả (Kaggle metric tính weighted F1)\n\nfor i, pid in enumerate(tqdm(test_ids_full)):\n    # Dùng numpy where để lọc nhanh\n    idxs = np.where(final_preds[i] > threshold)[0]\n    for idx in idxs:\n        term = top_terms[idx]\n        score = final_preds[i, idx]\n        submission_rows.append(f\"{pid}\\t{term}\\t{score:.3f}\")\n\nwith open('submission.tsv', 'w') as f:\n    # f.write(\"ProteinID\\tGO_ID\\tScore\\n\") # Bỏ comment nếu BTC yêu cầu header\n    f.write('\\n'.join(submission_rows))\n\nprint(\"DONE! File submission.tsv đã sẵn sàng trong tab Output.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}