{"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":14115963,"sourceType":"datasetVersion","datasetId":8992374}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Cell 1: Import & Config\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nfrom sklearn.model_selection import train_test_split\nimport gc\n\n# Cấu hình thiết bị\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {DEVICE}\")\n\n# Cấu hình đường dẫn (Dựa trên hình ảnh bạn cung cấp)\n# Lưu ý: Dataset tên là \"cafa6-t5-embeddings\" nhưng tên file lại là \"esm2\". \n# Tôi sẽ dùng đúng tên file trong hình.\nINPUT_EMBED_DIR = \"/kaggle/input/cafa6-t5-embeddings\"\n\nCONFIG = {\n    \"num_labels\": 1500,\n    \"batch_size\": 128,\n    \"lr\": 0.001,\n    \"epochs\": 20, # Tăng epoch lên vì chạy rất nhanh\n    \"paths\": {\n        # File Embedding có sẵn\n        \"train_embeds\": f\"{INPUT_EMBED_DIR}/train_embeddings_esm2.npy\",\n        \"train_ids\":    f\"{INPUT_EMBED_DIR}/train_ids_esm2.npy\",\n        \"test_embeds\":  f\"{INPUT_EMBED_DIR}/test_embeddings_esm2.npy\",\n        \"test_ids\":     f\"{INPUT_EMBED_DIR}/test_ids_esm2.npy\",\n        \n        # File Labels gốc của cuộc thi\n        \"train_terms\": \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n    }\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T13:35:46.871015Z","iopub.execute_input":"2025-12-14T13:35:46.871624Z","iopub.status.idle":"2025-12-14T13:35:52.323682Z","shell.execute_reply.started":"2025-12-14T13:35:46.871597Z","shell.execute_reply":"2025-12-14T13:35:52.322960Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 2: Load Embeddings\nprint(\"Loading embeddings from disk...\")\n\n# Load Train\ntrain_embeds = np.load(CONFIG['paths']['train_embeds'])\ntrain_ids = np.load(CONFIG['paths']['train_ids'], allow_pickle=True)\n\n# Load Test\ntest_embeds = np.load(CONFIG['paths']['test_embeds'])\ntest_ids = np.load(CONFIG['paths']['test_ids'], allow_pickle=True)\n\n# Tự động cập nhật kích thước đầu vào (Input Dimension)\n# Giúp code chạy đúng dù vector là T5 (1024) hay ESM2 (1280/2560)\nEMBED_DIM = train_embeds.shape[1]\n\nprint(f\"--- Data Loaded ---\")\nprint(f\"Train shape: {train_embeds.shape}\")\nprint(f\"Test shape:  {test_embeds.shape}\")\nprint(f\"Embedding Dimension detected: {EMBED_DIM}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T13:35:52.325002Z","iopub.execute_input":"2025-12-14T13:35:52.325404Z","iopub.status.idle":"2025-12-14T13:36:00.077108Z","shell.execute_reply.started":"2025-12-14T13:35:52.325385Z","shell.execute_reply":"2025-12-14T13:36:00.076392Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 3: Prepare Targets (SỬA LỖI ĐỊNH DẠNG ID: sp|ID|NAME -> ID)\nimport gc\n\nprint(\"--- Processing Targets & Fixing ID Format ---\")\n\n# 1. Load dữ liệu\ntrain_terms = pd.read_csv(CONFIG['paths']['train_terms'], sep=\"\\t\")\ntrain_ids = np.load(CONFIG['paths']['train_ids'], allow_pickle=True)\n\n# --- SỬA LỖI LOGIC TẠI ĐÂY ---\ndef clean_id(pid):\n    # Chuyển bytes sang string nếu cần\n    if isinstance(pid, bytes):\n        pid = pid.decode('utf-8')\n    pid_str = str(pid).strip()\n    \n    # Logic tách chuỗi: \"sp|A0A0C5B5G6|MOTSC_HUMAN\" -> lấy \"A0A0C5B5G6\"\n    # Tách bằng dấu gạch đứng '|' và lấy phần tử thứ 2 (index 1)\n    parts = pid_str.split('|')\n    if len(parts) > 1:\n        return parts[1] # Lấy mã ở giữa\n    return pid_str # Trả về nguyên gốc nếu không tìm thấy dấu |\n\n# Áp dụng hàm sửa lỗi\ntrain_ids_clean = [clean_id(pid) for pid in train_ids]\n\nprint(f\"[DEBUG] ID gốc: {train_ids[0]}\")\nprint(f\"[DEBUG] ID sau khi sửa: {train_ids_clean[0]}\") # Mong đợi: A0A0C5B5G6\n\n# Tạo map\nid_map = {pid: i for i, pid in enumerate(train_ids_clean)}\n\n# 2. Chọn Top Labels\ntop_terms = train_terms['term'].value_counts().index[:CONFIG['num_labels']]\nterm_to_idx = {term: i for i, term in enumerate(top_terms)}\n\n# 3. Tạo ma trận Targets\nnum_samples = len(train_ids)\nlabels_matrix = np.zeros((num_samples, CONFIG['num_labels']), dtype=np.float32)\n\n# Lọc dữ liệu\ntrain_terms['EntryID'] = train_terms['EntryID'].astype(str).str.strip()\nfiltered_terms = train_terms[\n    (train_terms['EntryID'].isin(id_map)) & \n    (train_terms['term'].isin(top_terms))\n]\n\nprint(f\"Số lượng dòng khớp được: {len(filtered_terms)}\")\n\nif len(filtered_terms) == 0:\n    print(\"❌ VẪN LỖI. Vui lòng kiểm tra lại logic cắt chuỗi.\")\nelse:\n    # Điền số 1 vào ma trận\n    print(\"Đang tạo ma trận nhãn...\")\n    for pid, term in tqdm(zip(filtered_terms['EntryID'], filtered_terms['term']), total=len(filtered_terms)):\n        row_idx = id_map[pid]\n        col_idx = term_to_idx[term]\n        labels_matrix[row_idx, col_idx] = 1.0\n        \n    print(f\"✅ THÀNH CÔNG! Tổng số nhãn dương: {labels_matrix.sum()}\")\n\n# Dọn dẹp\ndel train_terms, filtered_terms, train_ids_clean\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T13:36:00.077899Z","iopub.execute_input":"2025-12-14T13:36:00.078120Z","iopub.status.idle":"2025-12-14T13:36:01.393056Z","shell.execute_reply.started":"2025-12-14T13:36:00.078103Z","shell.execute_reply":"2025-12-14T13:36:01.392330Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 4: SE-ResNet Model Architecture (Advanced)\nclass ProteinDataset(Dataset):\n    def __init__(self, embeddings, targets=None, ids=None):\n        self.embeddings = embeddings\n        self.targets = targets\n        self.ids = ids\n        \n    def __len__(self):\n        return len(self.embeddings)\n    \n    def __getitem__(self, idx):\n        embed = torch.tensor(self.embeddings[idx], dtype=torch.float32)\n        if self.targets is not None:\n            target = torch.tensor(self.targets[idx], dtype=torch.float32)\n            return embed, target\n        return embed, self.ids[idx]\n\n# --- 1. Squeeze-and-Excitation Block (Cơ chế chú ý vào kênh quan trọng) ---\nclass SEBlock(nn.Module):\n    def __init__(self, channel, reduction=16):\n        super(SEBlock, self).__init__()\n        self.avg_pool = nn.AdaptiveAvgPool1d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(channel, channel // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(channel // reduction, channel, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x):\n        # x shape: [batch, features] -> cần unsqueeze để dùng avg_pool giả lập\n        b, c = x.size()\n        y = x.view(b, c, 1)\n        y = self.avg_pool(y).view(b, c)\n        y = self.fc(y)\n        return x * y # Scale lại đặc trưng gốc theo độ quan trọng\n\n# --- 2. Residual Block tích hợp SE ---\nclass AdvancedResBlock(nn.Module):\n    def __init__(self, in_features, out_features, dropout_rate=0.25):\n        super(AdvancedResBlock, self).__init__()\n        \n        self.block = nn.Sequential(\n            nn.Linear(in_features, out_features),\n            nn.BatchNorm1d(out_features),\n            nn.Mish(), # Mish activation tốt hơn ReLU trong mạng sâu\n            nn.Dropout(dropout_rate),\n            \n            nn.Linear(out_features, out_features),\n            nn.BatchNorm1d(out_features),\n            nn.Mish(),\n            nn.Dropout(dropout_rate)\n        )\n        \n        self.se = SEBlock(out_features) # Thêm SE Block vào cuối\n        \n        # Skip connection\n        self.shortcut = nn.Sequential()\n        if in_features != out_features:\n            self.shortcut = nn.Sequential(\n                nn.Linear(in_features, out_features),\n                nn.BatchNorm1d(out_features)\n            )\n\n    def forward(self, x):\n        identity = self.shortcut(x)\n        out = self.block(x)\n        out = self.se(out) # Áp dụng SE Attention\n        return out + identity\n\n# --- 3. Main Architecture: Deep Funnel ---\nclass SEResNet(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super(SEResNet, self).__init__()\n        \n        # Input Projection (Mở rộng chiều)\n        self.stem = nn.Sequential(\n            nn.Linear(input_dim, 2048),\n            nn.BatchNorm1d(2048),\n            nn.Mish(),\n            nn.Dropout(0.2)\n        )\n        \n        # Các lớp sâu hơn (Hình phễu: 2048 -> 1024 -> 512)\n        # Giúp nén thông tin, loại bỏ nhiễu\n        self.layer1 = AdvancedResBlock(2048, 2048, dropout_rate=0.3)\n        self.layer2 = AdvancedResBlock(2048, 1024, dropout_rate=0.3)\n        self.layer3 = AdvancedResBlock(1024, 1024, dropout_rate=0.3)\n        self.layer4 = AdvancedResBlock(1024, 512,  dropout_rate=0.3)\n        \n        # Classifier\n        self.head = nn.Sequential(\n            nn.Linear(512, num_classes)\n        )\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        x = self.layer4(x)\n        x = self.head(x)\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T13:36:01.393826Z","iopub.execute_input":"2025-12-14T13:36:01.394012Z","iopub.status.idle":"2025-12-14T13:36:01.405362Z","shell.execute_reply.started":"2025-12-14T13:36:01.393997Z","shell.execute_reply":"2025-12-14T13:36:01.404643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 5: Training (5-Fold + Focal Loss + Warm Restarts)\nfrom sklearn.model_selection import KFold\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n\n# --- CẤU HÌNH MỚI ---\nN_FOLDS = 5\nEPOCHS_PER_FOLD = 15 # Tăng epoch vì model sâu hơn cần train lâu hơn\nLR = 0.001\nBATCH_SIZE = 256\n\n# Định nghĩa Focal Loss (Tốt cho dữ liệu thưa như CAFA)\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=1, gamma=2, logits=True, reduce=True):\n        super(FocalLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.logits = logits\n        self.reduce = reduce\n        self.bce = nn.BCEWithLogitsLoss(reduction='none')\n\n    def forward(self, inputs, targets):\n        if self.logits:\n            bce_loss = self.bce(inputs, targets)\n        else:\n            bce_loss = nn.functional.binary_cross_entropy(inputs, targets, reduction='none')\n        pt = torch.exp(-bce_loss)\n        F_loss = self.alpha * (1-pt)**self.gamma * bce_loss\n\n        if self.reduce:\n            return torch.mean(F_loss)\n        else:\n            return F_loss\n\nkfold = KFold(n_splits=N_FOLDS, shuffle=True, random_state=42)\n\nprint(f\"\\n--- START TRAINING DEEP SE-RESNET ({N_FOLDS} FOLDS) ---\")\n\nfold_metrics = []\n\nfor fold, (train_idx, val_idx) in enumerate(kfold.split(train_embeds, labels_matrix)):\n    print(f\"\\n>>> FOLD {fold+1}/{N_FOLDS}\")\n    \n    # Prepare Data\n    X_train, X_val = train_embeds[train_idx], train_embeds[val_idx]\n    y_train, y_val = labels_matrix[train_idx], labels_matrix[val_idx]\n    \n    train_dataset = ProteinDataset(X_train, y_train)\n    val_dataset = ProteinDataset(X_val, y_val)\n    \n    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True)\n    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, pin_memory=True)\n    \n    # Init Model\n    model = SEResNet(input_dim=EMBED_DIM, num_classes=CONFIG['num_labels']).to(DEVICE)\n    \n    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-2) # Tăng weight decay để chống overfitting\n    \n    # Warm Restarts: Reset LR về max sau mỗi 5 epoch\n    scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=1, eta_min=1e-6)\n    \n    criterion = FocalLoss(gamma=2.0) # Dùng Focal Loss\n    \n    best_loss = float('inf')\n    \n    # Training Loop\n    for epoch in range(EPOCHS_PER_FOLD):\n        model.train()\n        train_loss = 0\n        for embeds, targets in tqdm(train_loader, desc=f\"Ep {epoch+1}\", leave=False):\n            embeds, targets = embeds.to(DEVICE), targets.to(DEVICE)\n            \n            optimizer.zero_grad()\n            outputs = model(embeds)\n            loss = criterion(outputs, targets)\n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item()\n        \n        # Validation\n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for embeds, targets in val_loader:\n                embeds, targets = embeds.to(DEVICE), targets.to(DEVICE)\n                outputs = model(embeds)\n                loss = criterion(outputs, targets) # Val loss cũng tính bằng Focal\n                val_loss += loss.item()\n        \n        avg_train_loss = train_loss / len(train_loader)\n        avg_val_loss = val_loss / len(val_loader)\n        \n        # Step scheduler theo từng epoch (cho WarmRestarts)\n        scheduler.step()\n        \n        if avg_val_loss < best_loss:\n            best_loss = avg_val_loss\n            torch.save(model.state_dict(), f\"model_fold_{fold}.pth\")\n            \n    print(f\"Fold {fold+1} Best Focal Loss: {best_loss:.5f}\")\n    fold_metrics.append(best_loss)\n\nprint(f\"\\n--- TRAINING FINISHED ---\")\nprint(f\"Avg Loss: {np.mean(fold_metrics):.5f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T13:36:01.406911Z","iopub.execute_input":"2025-12-14T13:36:01.407120Z","iopub.status.idle":"2025-12-14T13:44:19.751771Z","shell.execute_reply.started":"2025-12-14T13:36:01.407104Z","shell.execute_reply":"2025-12-14T13:44:19.751081Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 6: Ensemble Prediction (RAM Optimized - Incremental Write)\nimport gc\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\nprint(\"\\n--- Generating Ensemble Predictions (RAM Optimized) ---\")\n\n# 1. Khôi phục mapping\nprint(\"Re-creating index mapping...\")\ndf_terms = pd.read_csv(CONFIG['paths']['train_terms'], sep=\"\\t\")\ntop_terms = df_terms['term'].value_counts().index[:CONFIG['num_labels']]\nidx_to_term = {i: term for i, term in enumerate(top_terms)}\ndel df_terms, top_terms\ngc.collect()\n\n# 2. Setup Data Loader\n# Giảm Batch Size xuống chút để an toàn cho cả VRAM GPU\nBATCH_SIZE = 256 \ntest_dataset_final = ProteinDataset(test_embeds, ids=test_ids)\ntest_loader_final = DataLoader(test_dataset_final, batch_size=BATCH_SIZE, shuffle=False)\n\n# 3. Load Models\nmodels = []\nfor fold in range(N_FOLDS):\n    # Đảm bảo class SEResNet đã được định nghĩa ở Cell 4\n    model = SEResNet(input_dim=EMBED_DIM, num_classes=CONFIG['num_labels']).to(DEVICE)\n    model.load_state_dict(torch.load(f\"model_fold_{fold}.pth\"))\n    model.eval()\n    models.append(model)\n\nprint(f\"Loaded {len(models)} SE-ResNet models.\")\n\n# 4. Prediction Loop với cơ chế GHI FILE TRỰC TIẾP\nTHRESHOLD = 0.01\noutput_file = \"submission.tsv\"\n\nprint(f\"Starting inference... Writing directly to {output_file}\")\n\n# Mở file để ghi ngay lập tức ('w' mode)\nwith open(output_file, 'w') as f:\n    with torch.no_grad():\n        for step, (embeds, batch_ids) in enumerate(tqdm(test_loader_final)):\n            embeds = embeds.to(DEVICE)\n            \n            # Ensemble Averaging\n            avg_probs = None\n            for model in models:\n                logits = model(embeds)\n                probs = torch.sigmoid(logits)\n                if avg_probs is None:\n                    avg_probs = probs\n                else:\n                    avg_probs += probs\n            \n            avg_probs /= len(models)\n            avg_probs = avg_probs.cpu().numpy()\n            \n            # Xử lý kết quả của batch này\n            batch_lines = []\n            for i, pid in enumerate(batch_ids):\n                # Lọc ngưỡng\n                idx_above_thresh = np.where(avg_probs[i] > THRESHOLD)[0]\n                \n                for idx in idx_above_thresh:\n                    term = idx_to_term[idx]\n                    score = avg_probs[i][idx]\n                    # Tạo chuỗi line để ghi vào file (Id <tab> Term <tab> Score)\n                    batch_lines.append(f\"{pid}\\t{term}\\t{score:.3f}\\n\")\n            \n            # GHI NGAY XUỐNG Ổ CỨNG (Giải phóng RAM)\n            f.writelines(batch_lines)\n            \n            # Dọn dẹp biến tạm\n            del batch_lines, avg_probs, embeds\n            \n            # Gọi gom rác mỗi 50 steps để chắc chắn RAM được giải phóng\n            if step % 50 == 0:\n                gc.collect()\n\nprint(f\"✅ Submission saved to {output_file} successfully!\")\n\n# Kiểm tra vài dòng đầu của file\nprint(\"Sample of generated file:\")\nwith open(output_file, 'r') as f:\n    for _ in range(5):\n        print(f.readline().strip())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-14T13:44:19.752611Z","iopub.execute_input":"2025-12-14T13:44:19.753040Z","iopub.status.idle":"2025-12-14T13:50:59.133030Z","shell.execute_reply.started":"2025-12-14T13:44:19.753017Z","shell.execute_reply":"2025-12-14T13:50:59.132337Z"}},"outputs":[],"execution_count":null}]}