{"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,"isSourceIdPinned":false}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install mlflow\n!pip install lightning","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T12:00:49.570086Z","iopub.execute_input":"2025-12-13T12:00:49.570872Z","iopub.status.idle":"2025-12-13T12:01:03.273399Z","shell.execute_reply.started":"2025-12-13T12:00:49.570836Z","shell.execute_reply":"2025-12-13T12:01:03.272641Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q biopython\n\nimport os\nimport pandas as pd\nimport gc\nfrom Bio import SeqIO\n\n# 1. TẠO SEQUENCE.PARQUET\nseq_path = \"sequence.parquet\" \nfasta_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\n\nif not os.path.exists(seq_path):\n    print(f\"Creating {seq_path} from {fasta_path}...\")\n    seqs = []\n    if os.path.exists(fasta_path):\n        for record in SeqIO.parse(fasta_path, 'fasta'):\n            entry_id = record.id.split('|')[1] if '|' in record.id else record.id\n            seqs.append({'EntryID': entry_id})\n            \n            if len(seqs) % 50000 == 0:\n                print(f\"   Processed {len(seqs)} sequences...\")\n        \n        df_seq = pd.DataFrame(seqs)\n        df_seq.to_parquet(seq_path, index=False)\n        print(f\"Done! Created {seq_path} with shape {df_seq.shape}\")\n        del seqs, df_seq\n        gc.collect()\n    else:\n        print(f\"ERROR: Không tìm thấy file gốc tại {fasta_path}. Hãy kiểm tra lại dataset competition.\")\nelse:\n    print(f\"{seq_path} already exists.\")\n\nlabel_path = \"label_p.parquet\"\nterms_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n\nif not os.path.exists(label_path):\n    print(f\"Creating {label_path} from {terms_path}...\")\n    if os.path.exists(terms_path):\n        terms = pd.read_csv(terms_path, sep='\\t', usecols=['EntryID', 'term'])\n        top_k = 1000\n        top_terms = terms['term'].value_counts().head(top_k).index\n        terms_filtered = terms[terms['term'].isin(top_terms)]\n        \n        print(f\"   Generating label matrix for top {top_k} terms...\")\n        label_matrix = pd.get_dummies(terms_filtered, columns=['term'], prefix='term', sparse=True)\n        label_matrix = label_matrix.groupby('EntryID').max()\n        \n        label_df = label_matrix.reset_index()\n        \n        label_df.to_parquet(label_path, index=False)\n        print(f\"Done! Created {label_path} with shape {label_df.shape}\")\n        \n        del terms, terms_filtered, label_matrix, label_df\n        gc.collect()\n    else:\n         print(f\"ERROR: Không tìm thấy file gốc tại {terms_path}\")\nelse:\n    print(f\"{label_path} already exists.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T11:47:54.600802Z","iopub.execute_input":"2025-12-13T11:47:54.601445Z","iopub.status.idle":"2025-12-13T11:47:57.759991Z","shell.execute_reply.started":"2025-12-13T11:47:54.601416Z","shell.execute_reply":"2025-12-13T11:47:57.759061Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile config.py\nfrom types import SimpleNamespace\n\n# Cấu hình cho toàn bộ thí nghiệm\ncfg = SimpleNamespace( \n    learning_rate = 1e-3,\n    batch_size = 64,    \n    num_workers = 4, \n    # LƯU Ý: File của bạn tên là esm2.npy. ESM2 thường có dim là 1280 hoặc 640. \n    # T5 thường là 1024. Hãy chỉnh số này nếu chạy bị lỗi dimension mismatch.\n    input_dim = 1280,      \n    taxonomy_dim = 64,     \n    hidden_dim = 512,      \n    dropout = 0.1,\n    experiment_name = \"CAFA6_MultiModal_FixPath\",\n    tracking_uri = \"http://112.137.129.161:8836\", \n    num_epochs = 100,     \n    patience = 5, \n    loss = 'weighted_focal', \n    gamma = 2.0, \n\n    # --- ĐƯỜNG DẪN DỮ LIỆU ĐÃ SỬA THEO ẢNH ---\n    sequence_path = \"sequence.parquet\", \n    label_path = \"label_p.parquet\",\n    \n    # Đường dẫn taxonomy gốc của competition\n    train_taxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\", \n    \n    # Đường dẫn embedding (Lấy theo file trong ảnh dataset của bạn)\n    embedding_path = \"/kaggle/input/cafa6-t5-embeddings/train_embeddings_esm2.npy\", \n    \n    # Đường dẫn file IA\n    ia_path = \"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\", \n    \n    # Đường dẫn dữ liệu Test (Cần tạo hoặc trỏ đúng file nếu có)\n    # Hiện tại trong ảnh chưa thấy file test_protein.csv, mình để tạm placeholder\n    test_path = None, \n    test_taxonomy_path = None,\n    test_embedding_path = \"/kaggle/input/cafa6-t5-embeddings/test_embeddings_esm2.npy\", \n    \n    top_k = 1000, \n    output = \"submission_mlp_p.parquet\"\n)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-13T11:47:57.761228Z","iopub.execute_input":"2025-12-13T11:47:57.761538Z","iopub.status.idle":"2025-12-13T11:47:57.767869Z","shell.execute_reply.started":"2025-12-13T11:47:57.761505Z","shell.execute_reply":"2025-12-13T11:47:57.767138Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile metrics.py\nimport torch\nfrom torchmetrics import Metric\nimport numpy as np\n\nclass CAFAGenericMetric(Metric):\n    \"\"\"\n    A generic metric for CAFA-style evaluation. It calculates the maximum F-score (F-max)\n    over a range of thresholds. It can compute both standard F-max and a weighted F-max\n    if Information Content (IA) values are provided.\n\n    This metric is designed to work with multi-label classification problems where\n    predictions are continuous scores between 0 and 1.\n    \"\"\"\n    def __init__(self, ia: torch.Tensor = None, dist_sync_on_step=False):\n        \"\"\"\n        Args:\n            ia (torch.Tensor, optional): A 1D tensor containing the Information Content (IA)\n                value for each label. If provided, the metric will compute a weighted F-max.\n                The IA values are typically derived from an ontology file (e.g., .obo).\n                Defaults to None, which results in an unweighted F-max calculation.\n            dist_sync_on_step (bool): Synchronize metric state across processes at each step.\n                Defaults to False.\n        \"\"\"\n        super().__init__(dist_sync_on_step=dist_sync_on_step)\n        # Register state for predictions and targets.\n        # 'dist_reduce_fx=\"cat\"' ensures that when using distributed training,\n        # predictions and targets from all processes are concatenated.\n        self.add_state(\"preds\", default=[], dist_reduce_fx=\"cat\")\n        self.add_state(\"targets\", default=[], dist_reduce_fx=\"cat\")\n        \n        self.ia = ia\n\n    def update(self, preds: torch.Tensor, target: torch.Tensor):\n        \"\"\"\n        Update state with predictions and targets from a single batch.\n\n        Args:\n            preds: A tensor of predictions (scores), e.g., shape (batch_size, num_labels).\n            target: A tensor of ground truth labels (binary), e.g., shape (batch_size, num_labels).\n        \"\"\"\n        self.preds.append(preds)\n        self.targets.append(target)\n\n    def compute(self):\n        \"\"\"\n        Compute the F-max score over all accumulated predictions and targets.\n        It iterates through thresholds from 0.01 to 0.99 to find the threshold\n        that yields the highest F1-score. If IA values were provided, the calculation\n        is weighted.\n\n        Returns:\n            A tuple containing:\n            - fmax (float): The maximum F1-score (weighted or unweighted).\n            - precision (float): The precision at the threshold that gives F-max.\n            - recall (float): The recall at the threshold that gives F-max.\n        \"\"\"\n        preds = torch.cat(self.preds, dim=0).cpu().numpy()\n        targets = torch.cat(self.targets, dim=0).cpu().numpy()\n        \n        ia = self.ia.cpu().numpy() if self.ia is not None else None\n        if ia is not None:\n            # Ensure IA vector matches the number of labels\n            if ia.shape[0] != preds.shape[1]:\n                raise ValueError(\"The length of the IA vector must match the number of labels.\")\n\n        fmax = 0.0\n        prec = 0.0\n        rec = 0.0\n\n        # Iterate over thresholds\n        for t in range(1, 100):\n            tau = t / 100.0\n            \n            # Binarize predictions based on the current threshold\n            p = preds >= tau\n\n            if ia is None:\n                # --- Unweighted calculation ---\n                tp = np.sum(np.logical_and(p, targets))\n                fp = np.sum(np.logical_and(p, np.logical_not(targets)))\n                fn = np.sum(np.logical_and(np.logical_not(p), targets))\n            else:\n                # --- Weighted calculation ---\n                tp = np.sum(np.logical_and(p, targets) * ia)\n                fp = np.sum(np.logical_and(p, np.logical_not(targets)) * ia)\n                fn = np.sum(np.logical_and(np.logical_not(p), targets) * ia)\n\n            # Calculate precision\n            if tp + fp > 0:\n                precision = tp / (tp + fp)\n            else:\n                precision = 0\n\n            # Calculate recall\n            if tp + fn > 0:\n                recall = tp / (tp + fn)\n            else:\n                recall = 0\n\n            # Calculate F1-score\n            if precision + recall > 0:\n                f = 2 * (precision * recall) / (precision + recall)\n                # Update F-max if a better F1-score is found\n                if f > fmax:\n                    fmax = f\n                    prec = precision\n                    rec = recall\n        \n        return fmax, prec, rec\n\nclass CAFAMetrics(CAFAGenericMetric):\n    \"\"\"\n    CAFA Metric calculator using torchmetrics.\n    \"\"\"\n    def __init__(self, ia: torch.Tensor = None, dist_sync_on_step=False):\n        super().__init__(ia=ia, dist_sync_on_step=dist_sync_on_step)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T11:47:57.769382Z","iopub.execute_input":"2025-12-13T11:47:57.769556Z","iopub.status.idle":"2025-12-13T11:47:57.791194Z","shell.execute_reply.started":"2025-12-13T11:47:57.769543Z","shell.execute_reply":"2025-12-13T11:47:57.790572Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile losses.py\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.optim import AdamW\nimport lightning as L\nimport torchmetrics\n\nclass FocalLoss(nn.Module):\n    def __init__(self,\n                 alpha=None,  # có thể là float, list hoặc tensor\n                 gamma: float = 2.0,\n                 reduction: str = 'mean'):\n        super(FocalLoss, self).__init__()\n        if alpha is None:\n            self.alpha = None\n        elif isinstance(alpha, (list, tuple)):\n            self.alpha = torch.tensor(alpha, dtype=torch.float32)\n        elif isinstance(alpha, torch.Tensor):\n            self.alpha = alpha.float()\n        else:\n            self.alpha = torch.tensor([alpha], dtype=torch.float32)\n        self.gamma = gamma\n        self.reduction = reduction\n\n    def forward(self, inputs, targets):\n        \"\"\"\n        inputs: logits đầu ra (batch_size, num_classes)\n        targets: nhãn thật (batch_size, num_classes)\n        \"\"\"\n        BCE_loss = F.binary_cross_entropy_with_logits(inputs,\n                                                      targets,\n                                                      reduction='none')\n        pt = torch.exp(-BCE_loss)\n\n        if self.alpha is not None:\n            alpha = self.alpha.to(inputs.device)\n            if alpha.ndim == 1 and alpha.shape[0] == inputs.shape[1]:\n                alpha = alpha.unsqueeze(0)  # (1, num_classes)\n            F_loss = alpha * (1 - pt) ** self.gamma * BCE_loss\n        else:\n            F_loss = (1 - pt) ** self.gamma * BCE_loss\n\n        if self.reduction == 'mean':\n            return torch.mean(F_loss)\n        elif self.reduction == 'sum':\n            return torch.sum(F_loss)\n        else:\n            return F_loss\n\nclass WeightedFocalLoss(nn.Module):\n    def __init__(self,\n                 alpha=None,\n                 gamma: float = 2.0,\n                 term_weights: torch.Tensor = None):\n        \"\"\"\n        term_weights: Tensor shape (num_classes,) chứa trọng số IA cho từng label.\n        \"\"\"\n        super(WeightedFocalLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.term_weights = term_weights\n\n    def forward(self, inputs, targets):\n        \"\"\"\n        inputs: logits (B, C)\n        targets: labels (B, C)\n        \"\"\"\n       \n        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')\n        \n        pt = torch.exp(-bce_loss)\n        \n        \n        if self.alpha is not None:\n            # Alpha balancing (nếu dùng)\n            if isinstance(self.alpha, torch.Tensor):\n                alpha_t = self.alpha.to(inputs.device)\n            else:\n                alpha_t = self.alpha\n            focal_loss = alpha_t * (1 - pt) ** self.gamma * bce_loss\n        else:\n            focal_loss = (1 - pt) ** self.gamma * bce_loss\n            \n        # Áp dụng trọng số IA (Term Weights)\n        if self.term_weights is not None:\n            weights = self.term_weights.to(inputs.device)\n            # weights shape: (C,) -> broadcast thành (B, C)\n            weighted_loss = focal_loss * weights\n            return weighted_loss.sum() / (inputs.size(0) * weights.sum())\n            \n        else:\n            return focal_loss.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T11:47:57.791903Z","iopub.execute_input":"2025-12-13T11:47:57.792065Z","iopub.status.idle":"2025-12-13T11:47:57.809151Z","shell.execute_reply.started":"2025-12-13T11:47:57.792052Z","shell.execute_reply":"2025-12-13T11:47:57.808554Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile protein_module.py\nimport torch\nfrom torch import nn\nfrom torch.optim import AdamW\nimport lightning as L\nimport torchmetrics\nfrom losses import WeightedFocalLoss, FocalLoss\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom metrics import CAFAMetrics\n\nclass LitMLPModule(L.LightningModule):\n    def __init__(\n        self,\n        input_dim: int,\n        taxonomy_num: int,   # Tổng số lượng mã taxonomy\n        taxonomy_dim: int,   # Kích thước vector nhúng taxonomy\n        hidden_dim: int,\n        num_labels: int,\n        learning_rate: float,\n        term_weights: torch.Tensor, # Trọng số IA cho loss\n        cfg,\n    ):\n        super().__init__()\n        self.save_hyperparameters(ignore=['term_weights']) # Không save tensor lớn vào hparams\n        self.term_weights = term_weights\n        \n        self.tax_embedding = nn.Embedding(taxonomy_num, taxonomy_dim)\n        self.tax_head = nn.Sequential(\n            nn.Linear(taxonomy_dim, taxonomy_dim),\n            nn.ReLU()\n        )\n        \n        combined_dim = input_dim + taxonomy_dim\n        \n        self.main_mlp = nn.Sequential(\n            nn.Linear(combined_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(cfg.dropout),\n            \n            nn.Linear(hidden_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(cfg.dropout),\n            \n            nn.Linear(hidden_dim, num_labels) # Output layer (không sigmoid)\n        )\n        \n        # Loss function\n        if cfg.loss == 'focal':\n            self.criterion = FocalLoss(gamma=cfg.gamma)\n        elif cfg.loss == 'weighted_focal':\n            self.criterion = WeightedFocalLoss(gamma=cfg.gamma, term_weights=self.term_weights)\n        else: \n            self.criterion = nn.BCEWithLogitsLoss()\n\n        self.val_f1 = CAFAMetrics(ia=term_weights) \n\n    def forward(self, x_seq: torch.Tensor, x_tax: torch.Tensor):\n        # 1. Sequence Embedding\n        \n        # 2. Taxonomy Embedding\n        tax_emb = self.tax_embedding(x_tax)\n        tax_emb_processed = self.tax_head(tax_emb)\n        \n        # 3. Kết hợp\n        x_combined = torch.cat([x_seq, tax_emb_processed], dim=1)\n        \n        # 4. MLP\n        logits = self.main_mlp(x_combined)\n        return logits\n\n    def _common_step(self, batch, stage):\n        inputs = batch[\"embedding\"]\n        tax_ids = batch[\"taxonomy_idx\"]\n        labels = batch[\"labels\"]\n        \n        logits = self(inputs, tax_ids)\n        loss = self.criterion(logits, labels)\n        \n        preds = torch.sigmoid(logits) \n        \n        return loss, preds, labels\n\n    def training_step(self, batch, batch_idx):\n        loss, _, _ = self._common_step(batch, \"train\")\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        loss, preds, labels = self._common_step(batch, \"val\")\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True)\n        \n        self.val_f1.update(preds, labels.int())\n        \n        return loss\n    \n    def on_validation_epoch_end(self):\n        fmax, prec, rec = self.val_f1.compute()\n        self.log(\"val_fmax_w\", fmax, on_step=False, on_epoch=True, prog_bar=True)\n        self.val_f1.reset() # Reset state sau khi tính\n        \n    def predict_step(self, batch, batch_idx, dataloader_idx=0):\n        inputs = batch[\"embedding\"]\n        tax_ids = batch[\"taxonomy_idx\"]\n        protein_ids = batch[\"protein_id\"] \n        \n        logits = self(inputs, tax_ids)\n        preds = torch.sigmoid(logits)\n        return {\"predictions\": preds, \"ids\": protein_ids}\n\n    def configure_optimizers(self):\n        optimizer = AdamW(self.parameters(), lr=self.hparams.learning_rate)\n\n        # Scheduler theo dõi F-max (đã có trọng số IA)\n        scheduler = ReduceLROnPlateau(\n            optimizer,\n            mode='max',      # Maximize Val F-max weighted\n            factor=0.5,      # Reduce LR by half\n            patience=2,      # If 2 consecutive epochs F-max doesn't improve -> Reduce LR\n            min_lr=1e-6\n        )\n        \n        # Cấu hình để scheduler theo dõi 'val_fmax_w'\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"monitor\": \"val_fmax_w\",\n                \"interval\": \"epoch\",\n                \"frequency\": 1,\n            }\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T11:47:57.830664Z","iopub.execute_input":"2025-12-13T11:47:57.830837Z","iopub.status.idle":"2025-12-13T11:47:57.850386Z","shell.execute_reply.started":"2025-12-13T11:47:57.830823Z","shell.execute_reply":"2025-12-13T11:47:57.84965Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nimport lightning as L\nfrom sklearn.model_selection import train_test_split\nfrom lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping\n# Thay đổi Logger\nfrom lightning.pytorch.loggers import CSVLogger \nfrom torch.utils.data import Dataset, DataLoader\n# Import module nhưng ta sẽ định nghĩa đè class LitProteinDataModule bên dưới\nfrom protein_module import LitMLPModule\nfrom config import cfg \nimport gc\nimport os\nimport pyarrow as pa\nimport pyarrow.parquet as pq\n\nclass ProteinEmbeddingDataset(Dataset):\n    def __init__(self, embeddings, dataframe, label_columns=None, taxonomy_map=None):\n        self.embeddings = embeddings\n        self.dataframe = dataframe\n        self.label_columns = label_columns\n        self.taxonomy_map = taxonomy_map\n        \n        if 'EntryID' in self.dataframe.columns:\n            self.id_col = 'EntryID'\n        elif 'id' in self.dataframe.columns:\n            self.id_col = 'id'\n        else:\n            self.id_col = self.dataframe.columns[0]\n            \n        self.tax_col = 'taxonomyID'\n        self.unknown_tax_idx = 0 \n        if taxonomy_map is not None:\n             self.unknown_tax_idx = taxonomy_map.get('<unknown>', 0)\n\n    def __len__(self):\n        return len(self.dataframe)\n\n    def __getitem__(self, index: int):\n        emb_vector = self.embeddings[index]\n        data_row = self.dataframe.iloc[index]\n        protein_id = data_row[self.id_col]\n        \n        item = dict(\n            protein_id=protein_id, \n            embedding=torch.tensor(emb_vector, dtype=torch.float32)\n        )\n        \n        tax_id = self.unknown_tax_idx \n        if self.tax_col in data_row and not pd.isna(data_row[self.tax_col]) and self.taxonomy_map is not None:\n            tax_id = self.taxonomy_map.get(data_row[self.tax_col], self.unknown_tax_idx)\n        item['taxonomy_idx'] = torch.tensor(tax_id, dtype=torch.long)\n            \n        if self.label_columns is not None:\n            labels = data_row[self.label_columns].values.astype(np.float32)\n            item['labels'] = torch.tensor(labels, dtype=torch.float32)\n            \n        return item\n\nclass LitProteinDataModule(L.LightningDataModule):\n    def __init__(self, cfg, train_df, train_emb, label_cols, taxonomy_map, test_df=None, test_emb=None):\n        super().__init__()\n        self.cfg = cfg\n        self.label_columns = label_cols\n        self.taxonomy_map = taxonomy_map\n        \n        # --- LOGIC SPLIT ĐÚNG ---\n        train_split_df, val_split_df = train_test_split(\n            train_df,\n            test_size=0.1, \n            random_state=42, \n            shuffle=True\n        )\n        \n        t_idxs = train_split_df.index\n        v_idxs = val_split_df.index\n        \n        self.train_emb = train_emb[t_idxs]\n        self.val_emb = train_emb[v_idxs]\n        \n        self.train_df = train_split_df.reset_index(drop=True)\n        self.val_df = val_split_df.reset_index(drop=True)\n        \n        self.test_df = test_df\n        self.test_emb = test_emb\n        \n    def setup(self, stage=None):\n        self.train_dataset = ProteinEmbeddingDataset(\n            embeddings=self.train_emb,\n            dataframe=self.train_df, \n            label_columns=self.label_columns,\n            taxonomy_map=self.taxonomy_map\n        )\n        self.val_dataset = ProteinEmbeddingDataset(\n            embeddings=self.val_emb,\n            dataframe=self.val_df, \n            label_columns=self.label_columns,\n            taxonomy_map=self.taxonomy_map\n        )\n        if self.test_df is not None:\n            self.test_dataset = ProteinEmbeddingDataset(\n                embeddings=self.test_emb,\n                dataframe=self.test_df, \n                label_columns=None,\n                taxonomy_map=self.taxonomy_map\n            )\n\n    def train_dataloader(self):\n        return DataLoader(self.train_dataset, batch_size=self.cfg.batch_size, shuffle=True, num_workers=self.cfg.num_workers)\n\n    def val_dataloader(self):\n        return DataLoader(self.val_dataset, batch_size=self.cfg.batch_size, num_workers=self.cfg.num_workers)\n\n    def test_dataloader(self):\n        if self.test_dataset is None: return None\n        return DataLoader(self.test_dataset, batch_size=self.cfg.batch_size, num_workers=self.cfg.num_workers)\n\n# ==============================================================================\n# 2. CONFIG PATHS & MAIN\n# ==============================================================================\ncfg.train_taxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\"\ncfg.embedding_path = \"/kaggle/input/cafa6-t5-embeddings/train_embeddings_esm2.npy\"\ncfg.ia_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/IA.tsv\"\ncfg.input_dim = 1280 \n\ndef load_data():\n    print(\"1. Loading Training Data...\")\n    seq_df = pd.read_parquet(cfg.sequence_path)\n    \n    if os.path.exists(cfg.train_taxonomy_path):\n        tax_df = pd.read_csv(cfg.train_taxonomy_path, sep='\\t')\n        if 'EntryID' not in tax_df.columns:\n            tax_df.rename(columns={tax_df.columns[0]: 'EntryID', tax_df.columns[1]: 'taxonomyID'}, inplace=True)\n        seq_df = seq_df.merge(tax_df[['EntryID', 'taxonomyID']], on='EntryID', how='left')\n        seq_df['taxonomyID'] = seq_df['taxonomyID'].fillna(-1)\n    \n    print(f\"   Loading Embeddings from {cfg.embedding_path}...\")\n    embeddings = np.load(cfg.embedding_path)\n    \n    min_len = min(len(seq_df), len(embeddings))\n    seq_df = seq_df.iloc[:min_len]\n    embeddings = embeddings[:min_len]\n    \n    lbl_df = pd.read_parquet(cfg.label_path)\n    \n    print(\"2. Merging Data & Labels...\")\n    seq_df['original_index'] = seq_df.index\n    full_df = seq_df.merge(lbl_df, on='EntryID', how='inner')\n    \n    valid_indices = full_df['original_index'].values\n    filtered_embeddings = embeddings[valid_indices]\n    \n    full_df = full_df.reset_index(drop=True)\n    \n    exclude = ['EntryID', 'taxonomyID', 'original_index', 'seq']\n    label_cols = [c for c in full_df.columns if c in lbl_df.columns and c not in exclude]\n    \n    if cfg.top_k > 0:\n        counts = full_df[label_cols].sum().sort_values(ascending=False)\n        label_cols = counts.head(cfg.top_k).index.tolist()\n        \n    unique_tax = full_df['taxonomyID'].unique()\n    tax_map = {'<unknown>': 0}\n    curr = 1\n    for t in unique_tax:\n        if t != -1:\n            tax_map[t] = curr\n            curr += 1\n            \n    return full_df, filtered_embeddings, label_cols, tax_map, len(tax_map)\n\ndef load_ia(ia_path, label_cols):\n    if not os.path.exists(ia_path):\n        print(f\"   Warning: IA file not found at {ia_path}\")\n        return None\n    try:\n        ia_df = pd.read_csv(ia_path, sep='\\t', header=None, names=['term', 'score'])\n        w_dict = dict(zip(ia_df['term'], ia_df['score']))\n        weights = []\n        for col in label_cols:\n            term = col.replace('term_', '')\n            weights.append(w_dict.get(term, 1.0))\n        return torch.tensor(weights, dtype=torch.float32)\n    except:\n        return None\n\ndef main():\n    L.seed_everything(42)\n    \n    full_df, full_emb, label_cols, tax_map, num_tax = load_data()\n    term_weights = load_ia(cfg.ia_path, label_cols)\n    \n    dm = LitProteinDataModule(\n        cfg=cfg,\n        train_df=full_df,\n        train_emb=full_emb,\n        label_cols=label_cols,\n        taxonomy_map=tax_map\n    )\n    \n    print(f\"3. Init Model (In: {full_emb.shape[1]}, Tax: {num_tax}, Out: {len(label_cols)})\")\n    \n    model = LitMLPModule(\n        input_dim=full_emb.shape[1],\n        taxonomy_num=num_tax,\n        taxonomy_dim=cfg.taxonomy_dim,\n        hidden_dim=cfg.hidden_dim,\n        num_labels=len(label_cols),\n        learning_rate=cfg.learning_rate,\n        term_weights=term_weights,\n        cfg=cfg\n    )\n    \n    checkpoint = ModelCheckpoint(\n        dirpath=\"checkpoints\", filename=\"best-model\",\n        save_top_k=1, monitor=\"val_fmax_w\", mode=\"max\"\n    )\n    \n    csv_logger = CSVLogger(\"logs\", name=cfg.experiment_name)\n    \n    trainer = L.Trainer(\n        max_epochs=cfg.num_epochs,\n        callbacks=[checkpoint],\n        accelerator=\"auto\",\n        logger=csv_logger  \n    )\n    \n    print(\"Start Training...\")\n    trainer.fit(model, dm)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T12:02:29.568796Z","iopub.execute_input":"2025-12-13T12:02:29.569462Z","iopub.status.idle":"2025-12-13T13:28:57.383324Z","shell.execute_reply.started":"2025-12-13T12:02:29.569435Z","shell.execute_reply":"2025-12-13T13:28:57.382642Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom protein_module import LitMLPModule\nfrom config import cfg\nimport os\nimport types \nimport gc\n\ndef generate_submission():\n    print(\"--- STARTING INFERENCE ---\")\n    \n    # 1. Load Model\n    checkpoint_path = \"checkpoints/best-model.ckpt\"\n    if not os.path.exists(checkpoint_path):\n        print(\"Không tìm thấy checkpoint. Hãy đảm bảo quá trình train đã chạy xong.\")\n        return\n\n    print(f\"Loading model from {checkpoint_path}...\")\n    \n    torch.serialization.add_safe_globals([types.SimpleNamespace]) \n    \n    model = LitMLPModule.load_from_checkpoint(checkpoint_path, term_weights=None)\n    \n    model.eval()\n    model.cuda()\n    \n    # 2. Load Test Data\n    test_emb_path = \"/kaggle/input/cafa6-t5-embeddings/test_embeddings_esm2.npy\"\n    test_ids_path = \"/kaggle/input/cafa6-t5-embeddings/test_ids_esm2.npy\"\n    test_taxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\"\n    \n    print(f\"Loading Test Embeddings from {test_emb_path}...\")\n    test_emb = np.load(test_emb_path)\n    print(f\"Loading Test IDs from {test_ids_path}...\")\n    test_ids = np.load(test_ids_path)\n    \n    test_df = pd.DataFrame({'EntryID': test_ids})\n    \n    test_df['EntryID'] = test_df['EntryID'].astype(str)\n    \n    # Merge Taxonomy\n    if os.path.exists(test_taxonomy_path):\n        print(\"Merging Test Taxonomy...\")\n        test_tax_df = pd.read_csv(test_taxonomy_path, sep='\\t')\n        \n        if 'EntryID' not in test_tax_df.columns:\n    \n             test_tax_df.rename(columns={test_tax_df.columns[0]: 'EntryID', test_tax_df.columns[1]: 'taxonomyID'}, inplace=True)\n       \n        test_df = test_df.merge(test_tax_df[['EntryID', 'taxonomyID']], on='EntryID', how='left')\n        test_df['taxonomyID'] = test_df['taxonomyID'].fillna(-1).astype(int)\n    else:\n        test_df['taxonomyID'] = -1\n\n    print(\"Reconstructing Label Columns...\")\n    lbl_df_temp = pd.read_parquet(\"label_p.parquet\")\n    exclude_cols = ['EntryID', 'taxonomyID', 'original_index', 'seq']\n    all_label_cols = [c for c in lbl_df_temp.columns if c not in exclude_cols]\n    \n    if cfg.top_k is not None and cfg.top_k > 0:\n        label_counts = lbl_df_temp[all_label_cols].sum().sort_values(ascending=False)\n        label_cols = label_counts.head(cfg.top_k).index.tolist()\n    else:\n        label_cols = all_label_cols\n        \n    num_labels_model = model.main_mlp[-1].out_features\n    if len(label_cols) != num_labels_model:\n        print(f\"FATAL ERROR: Label count mismatch! Model: {num_labels_model}, Recalculated: {len(label_cols)}\")\n        return\n\n    # 4. Predict Loop\n    print(\"Using default taxonomy mapping (0) for safety...\")\n    batch_size = 512 \n    num_samples = len(test_df)\n    results = []\n    \n    print(f\"Predicting on {num_samples} samples...\")\n    \n    tax_indices_fixed = torch.zeros(batch_size, dtype=torch.long).cuda()\n    \n    with torch.no_grad():\n        for i in tqdm(range(0, num_samples, batch_size)):\n            end = min(i + batch_size, num_samples)\n            \n            # Batch Data\n            batch_emb = torch.tensor(test_emb[i:end], dtype=torch.float32).cuda()\n            batch_ids = test_ids[i:end]\n            \n            # Batch Tax\n            curr_bs = end - i\n            batch_tax = tax_indices_fixed[:curr_bs]\n            \n            # Forward\n            logits = model(batch_emb, batch_tax)\n            probs = torch.sigmoid(logits).cpu().numpy()\n            \n            # Process Output\n            for b_idx in range(curr_bs):\n                pid = batch_ids[b_idx]\n                p_scores = probs[b_idx]\n                \n                # Filter > 0.01\n                mask = p_scores > 0.01\n                if mask.any():\n                    valid_indices = np.where(mask)[0]\n                    for idx in valid_indices:\n                        term_name = label_cols[idx].replace('term_', '') \n                        score = p_scores[idx]\n                        results.append(f\"{pid}\\t{term_name}\\t{score:.3f}\")\n                        \n    # 5. Save Submission\n    out_file = \"submission.tsv\"\n    print(f\"Writing {len(results)} predictions to {out_file}...\")\n    \n    with open(out_file, \"w\") as f:\n        for line in results:\n            f.write(line + \"\\n\")\n            \n    print(\"Done! Submission saved.\")\n    gc.collect()\n\nif __name__ == \"__main__\":\n    generate_submission()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T13:38:10.662805Z","iopub.execute_input":"2025-12-13T13:38:10.66332Z","iopub.status.idle":"2025-12-13T13:39:57.852291Z","shell.execute_reply.started":"2025-12-13T13:38:10.663293Z","shell.execute_reply":"2025-12-13T13:39:57.85151Z"}},"outputs":[],"execution_count":null}]}