{"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":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14115963,"sourceType":"datasetVersion","datasetId":8992374}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q --upgrade pip\n!pip install -q lightning torchmetrics\n!pip install -q mlflow\n!pip install -q --no-cache-dir mlflow ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:45.660768Z","iopub.execute_input":"2025-12-12T00:32:45.661325Z","iopub.status.idle":"2025-12-12T00:32:57.177804Z","shell.execute_reply.started":"2025-12-12T00:32:45.661283Z","shell.execute_reply":"2025-12-12T00:32:57.176284Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import mlflow\nimport lightning as L\nfrom lightning.pytorch.loggers import MLFlowLogger\n\nprint(\"MLFlow version:\", mlflow.__version__)\nprint(\"Lightning imported successfully!\")\nprint(\"MLFlowLogger ready!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.180751Z","iopub.execute_input":"2025-12-12T00:32:57.181371Z","iopub.status.idle":"2025-12-12T00:32:57.187767Z","shell.execute_reply.started":"2025-12-12T00:32:57.181336Z","shell.execute_reply":"2025-12-12T00:32:57.186638Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 1: Imports and Setup\n# Combine all necessary imports from the files\n\nimport 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\nfrom lightning.pytorch.loggers import MLFlowLogger\nimport gc\nimport os\nfrom tqdm import tqdm\nimport pyarrow as pa\nimport pyarrow.parquet as pq\nimport types\nfrom types import SimpleNamespace\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import AdamW\nimport torchmetrics\nfrom torchmetrics import Metric\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\n# Optional: For safe serialization (from train.py)\n# torch.serialization.add_safe_globals([types.SimpleNamespace])","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.18894Z","iopub.execute_input":"2025-12-12T00:32:57.189196Z","iopub.status.idle":"2025-12-12T00:32:57.217275Z","shell.execute_reply.started":"2025-12-12T00:32:57.189177Z","shell.execute_reply":"2025-12-12T00:32:57.216058Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 2: Config (sửa lỗi cú pháp và đầy đủ thuộc tính)\ncfg = SimpleNamespace(\n    learning_rate = 1e-3,\n    batch_size = 64,\n    num_workers = 4,\n    input_dim = 1280,  # ESM2\n    taxonomy_dim = 64,\n    hidden_dim = 512,\n    dropout = 0.1,\n    experiment_name = \"CAFA6_ESM2_MultiModal\",\n    tracking_uri = None,  # Local MLFlow\n    num_epochs = 100,\n    patience = 2,\n    loss = 'bce',\n    gamma = 2.0,\n    \n    # Train data\n    sequence_path = \"/kaggle/working/sequence.parquet\",\n    train_taxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\",\n    label_path = \"/kaggle/working/label_p.parquet\",\n    ia_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/IA.tsv\",\n    \n    # Test data (dùng embeddings + IDs có sẵn)\n    test_embedding_path = \"/kaggle/input/cafa6-t5-embeddings/test_embeddings_esm2.npy\",\n    test_ids_path = \"/kaggle/input/cafa6-t5-embeddings/test_ids_esm2.npy\",  # File IDs\n    test_taxonomy_path = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\",  # Optional\n    \n    # Train embeddings\n    embedding_path = \"/kaggle/input/cafa6-t5-embeddings/train_embeddings_esm2.npy\",\n    \n    top_k = 1000,\n    output = \"submission.parquet\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.219875Z","iopub.execute_input":"2025-12-12T00:32:57.220179Z","iopub.status.idle":"2025-12-12T00:32:57.237516Z","shell.execute_reply.started":"2025-12-12T00:32:57.220155Z","shell.execute_reply":"2025-12-12T00:32:57.236535Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 3: Losses (from losses.py)\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        # Tính Binary Cross Entropy Loss không reduction\n        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')\n        \n        pt = torch.exp(-bce_loss)\n        \n        # Tính Focal term\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-12T00:32:57.23908Z","iopub.execute_input":"2025-12-12T00:32:57.239424Z","iopub.status.idle":"2025-12-12T00:32:57.270312Z","shell.execute_reply.started":"2025-12-12T00:32:57.239394Z","shell.execute_reply":"2025-12-12T00:32:57.269128Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 4: Metrics (from metrics.py)\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    This class inherits from CAFAGenericMetric and is intended to be used\n    within a PyTorch Lightning module or a standard PyTorch training loop.\n    It can compute both a standard F-max and a weighted F-max using\n    Information Content (IA) values.\n\n    Example (Unweighted):\n    ---------------------\n    >>> # Mocks for demonstration\n    >>> num_batches = 10\n    >>> num_classes = 1500\n    >>>\n    >>> # 1. Initialize the metric for unweighted F-max\n    >>> cafa_metric = CAFAMetrics()\n    >>>\n    >>> # 2. Simulate a validation loop\n    >>> for _ in range(num_batches):\n    ...     preds = torch.rand(32, num_classes)\n    ...     targets = torch.randint(0, 2, (32, num_classes))\n    ...     cafa_metric.update(preds, targets)\n    >>>\n    >>> # 3. Compute the final metric\n    >>> fmax, precision, recall = cafa_metric.compute()\n    >>> print(f\"Unweighted F-max: {fmax:.4f}\")\n\n    Example (Weighted):\n    -------------------\n    >>> # 1. Initialize the metric with IA values for weighted F-max\n    >>> # In a real scenario, IA values are derived from ontology files.\n    >>> ia_values = torch.rand(num_classes)\n    >>> cafa_metric_weighted = CAFAMetrics(ia=ia_values)\n    >>>\n    >>> # 2. Simulate a validation loop\n    >>> for _ in range(num_batches):\n    ...     preds = torch.rand(32, num_classes)\n    ...     targets = torch.randint(0, 2, (32, num_classes))\n    ...     cafa_metric_weighted.update(preds, targets)\n    >>>\n    >>> # 3. Compute the final metric\n    >>> fmax_w, _, _ = cafa_metric_weighted.compute()\n    >>> print(f\"Weighted F-max: {fmax_w:.4f}\")\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-12T00:32:57.271494Z","iopub.execute_input":"2025-12-12T00:32:57.271869Z","iopub.status.idle":"2025-12-12T00:32:57.298685Z","shell.execute_reply.started":"2025-12-12T00:32:57.271839Z","shell.execute_reply":"2025-12-12T00:32:57.297656Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 5: Protein Dataset (from protein_dataset.py)\n\nclass ProteinEmbeddingDataset(Dataset):\n    def __init__(self,\n                 embeddings: np.ndarray,\n                 dataframe: pd.DataFrame,\n                 label_columns: list = None,\n                 taxonomy_map: dict = None):\n        \n        self.embeddings = embeddings\n        self.dataframe = dataframe\n        self.label_columns = label_columns\n        self.taxonomy_map = taxonomy_map\n        \n        # Tự động tìm cột ID\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        # Kiểm tra cột Taxonomy\n        self.tax_col = 'taxonomyID'\n        # Default index cho unknown tax (thường là 0)\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        # 1. Sequence Embedding\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        # 2. Taxonomy Indexing\n        if self.taxonomy_map is not None and self.tax_col in data_row:\n            raw_tax = data_row[self.tax_col]\n            # Map raw tax ID sang index, nếu không có trong map thì về unknown\n            tax_idx = self.taxonomy_map.get(raw_tax, self.unknown_tax_idx)\n            item['taxonomy_idx'] = torch.tensor(tax_idx, dtype=torch.long)\n        else:\n            # Fallback an toàn\n            item['taxonomy_idx'] = torch.tensor(self.unknown_tax_idx, dtype=torch.long)\n        \n        # 3. Labels\n        if self.label_columns is not None:\n            labels = data_row[self.label_columns].values.astype(float)\n            item['labels'] = torch.FloatTensor(labels)\n            \n        return item\n\n\nclass LitProteinDataModule(L.LightningDataModule):\n    def __init__(self,\n                 train_df: pd.DataFrame,\n                 train_emb: np.ndarray,\n                 val_df: pd.DataFrame,\n                 val_emb: np.ndarray,\n                 label_columns: list,\n                 taxonomy_map: dict,\n                 test_df: pd.DataFrame = None,\n                 test_emb: np.ndarray = None,\n                 batch_size: int = 32,\n                 num_workers: int = 4):\n        super().__init__()\n        self.train_df = train_df\n        self.train_emb = train_emb\n        self.val_df = val_df\n        self.val_emb = val_emb\n        \n        self.test_df = test_df\n        self.test_emb = test_emb\n        \n        self.label_columns = label_columns\n        self.taxonomy_map = taxonomy_map\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        \n    def setup(self, stage: str = 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        \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        \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.batch_size, shuffle=True, num_workers=self.num_workers)\n\n    def val_dataloader(self):\n        return DataLoader(self.val_dataset, batch_size=self.batch_size, num_workers=self.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.batch_size, num_workers=self.num_workers, shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.29979Z","iopub.execute_input":"2025-12-12T00:32:57.300111Z","iopub.status.idle":"2025-12-12T00:32:57.330221Z","shell.execute_reply.started":"2025-12-12T00:32:57.30008Z","shell.execute_reply":"2025-12-12T00:32:57.32898Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 6: Protein Module (sửa BatchNorm → LayerNorm để tránh lỗi batch_size=1)\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        # 1. Taxonomy Branch\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        # Main MLP: Thay BatchNorm1d → LayerNorm (ổn định với batch_size=1)\n        self.main_mlp = nn.Sequential(\n            nn.Linear(combined_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),   # ← Thay đổi quan trọng\n            nn.ReLU(),\n            nn.Dropout(cfg.dropout),\n            \n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.LayerNorm(hidden_dim // 2),  # ← Thay đổi quan trọng\n            nn.ReLU(),\n            nn.Dropout(cfg.dropout),\n            \n            nn.Linear(hidden_dim // 2, num_labels)\n        )\n\n        # Loss Function\n        if cfg.loss == 'weighted_focal':\n            # self.criterion = WeightedFocalLoss(gamma=cfg.gamma, term_weights=self.term_weights)\n            self.criterion = FocalLoss(alpha=getattr(cfg, 'alpha', None), gamma=cfg.gamma)\n        else:\n            self.criterion = nn.BCEWithLogitsLoss()\n\n        # Validation Metrics\n        self.val_f1 = torchmetrics.F1Score(task=\"multilabel\", num_labels=num_labels, average='micro')\n\n    def forward(self, x_seq, x_tax):\n        tax_emb = self.tax_embedding(x_tax)  # (B, Tax_Dim)\n        tax_feat = self.tax_head(tax_emb)\n        \n        combined = torch.cat([x_seq, tax_feat], dim=1)\n        \n        # Main MLP\n        logits = self.main_mlp(combined)\n        return logits\n\n    def _common_step(self, batch, stage: str):\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        preds = torch.sigmoid(logits)\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(preds, labels.int())\n        self.log(\"val_f1\", self.val_f1, on_step=False, on_epoch=True, prog_bar=True)\n\n        return loss\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 = ReduceLROnPlateau(\n            optimizer,\n            mode='max',      # Maximize val_f1\n            factor=0.5,\n            patience=2,\n            min_lr=1e-6\n        )\n        \n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"monitor\": \"val_f1\",\n                \"interval\": \"epoch\",\n                \"frequency\": 1,\n            },\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.331408Z","iopub.execute_input":"2025-12-12T00:32:57.331744Z","iopub.status.idle":"2025-12-12T00:32:57.351028Z","shell.execute_reply.started":"2025-12-12T00:32:57.331718Z","shell.execute_reply":"2025-12-12T00:32:57.350061Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 7: Train Functions (sửa hoàn chỉnh cho CAFA 6 - không còn cột 'seq')\n\ndef load_and_align_train_data():\n    print(\"1. Loading Training Data...\")\n    \n    # A. Load Sequence (để lấy EntryID làm gốc)\n    print(f\"   Loading Sequence from {cfg.sequence_path}...\")\n    seq_df = pd.read_parquet(cfg.sequence_path)\n    \n    # B. Load Taxonomy Train (train_taxonomy.tsv - no header)\n    if os.path.exists(cfg.train_taxonomy_path):\n        print(f\"   Loading Train Taxonomy from {cfg.train_taxonomy_path}...\")\n        tax_df = pd.read_csv(cfg.train_taxonomy_path, sep='\\t', header=None, names=['EntryID', 'taxonomyID'])\n        \n        if 'EntryID' not in tax_df.columns or 'taxonomyID' not in tax_df.columns:\n            raise ValueError(f\"File taxonomy {cfg.train_taxonomy_path} phải có 2 cột (EntryID và taxonomyID)\")\n        \n        seq_df = seq_df.merge(tax_df[['EntryID', 'taxonomyID']], on='EntryID', how='left')\n        seq_df['taxonomyID'] = seq_df['taxonomyID'].fillna(-1).astype(int)\n    else:\n        raise FileNotFoundError(f\"Không tìm thấy file taxonomy train tại: {cfg.train_taxonomy_path}\")\n\n    # C. Load Embeddings\n    print(f\"   Loading embeddings from {cfg.embedding_path}...\")\n    embeddings = np.load(cfg.embedding_path)\n    \n    if len(seq_df) != len(embeddings):\n        raise ValueError(f\"Mismatch: Sequence DF ({len(seq_df)}) vs Embeddings ({len(embeddings)})\")\n\n    # D. Load Labels & Merge\n    print(f\"   Loading Labels from {cfg.label_path}...\")\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    \n    # Inner merge để chỉ giữ protein có nhãn\n    full_df = seq_df.merge(lbl_df, on='EntryID', how='inner')\n    \n    # Lọc embedding tương ứng\n    valid_indices = full_df['original_index'].values\n    filtered_embeddings = embeddings[valid_indices]\n    \n    print(f\"   Final Train Shape: DF {full_df.shape}, Emb {filtered_embeddings.shape}\")\n    \n    # Xác định label columns (loại bỏ metadata - KHÔNG có 'seq')\n    metadata_cols = ['EntryID', 'taxonomyID', 'original_index']\n    label_cols = [c for c in full_df.columns if c not in metadata_cols]\n    \n    if cfg.top_k is not None and cfg.top_k > 0:\n        print(f\"   Filtering for top {cfg.top_k} most frequent labels...\")\n        label_counts = full_df[label_cols].sum().sort_values(ascending=False)\n        top_k_labels = label_counts.head(cfg.top_k).index.tolist()\n        print(f\"   Original labels: {len(label_cols)} → After top-k: {len(top_k_labels)}\")\n        label_cols = top_k_labels\n        \n        # Chỉ giữ các cột cần thiết (không có 'seq')\n        full_df = full_df[metadata_cols + label_cols].copy()\n    else:\n        print(\"   Skipping top-k label filtering. Using all labels.\")\n\n    # E. Build Taxonomy Map\n    print(\"3. Building Taxonomy Map...\")\n    unique_taxonomies = full_df['taxonomyID'].unique()\n    taxonomy_map = {'<unknown>': 0}\n    current_idx = 1\n    for val in unique_taxonomies:\n        if val != -1:\n            taxonomy_map[val] = current_idx\n            current_idx += 1\n    num_taxonomies = len(taxonomy_map)\n    print(f\"   Found {num_taxonomies - 1} unique taxonomies (plus unknown).\")\n    \n    return full_df, filtered_embeddings, label_cols, taxonomy_map, num_taxonomies\n\n\ndef load_test_data(taxonomy_map):\n    print(\"   Loading Test Data from embeddings and IDs...\")\n    if not os.path.exists(cfg.test_embedding_path):\n        print(\"   Test embeddings not found. Skipping prediction.\")\n        return None, None\n\n    test_emb = np.load(cfg.test_embedding_path)\n    \n    # Load test IDs\n    ids_path = getattr(cfg, 'test_ids_path', cfg.test_embedding_path.replace('.npy', '_ids.npy'))\n    if os.path.exists(ids_path):\n        test_ids = np.load(ids_path)\n    else:\n        print(\"   Warning: test_ids not found. Using dummy IDs.\")\n        test_ids = np.arange(len(test_emb)).astype(str)\n    \n    # Dummy test_df (chỉ cần id và taxonomyID = -1)\n    test_df = pd.DataFrame({\n        'id': test_ids,\n        'taxonomyID': np.full(len(test_emb), -1, dtype=int)\n    })\n    \n    return test_df, test_emb\n\ndef load_ia_weights(ia_path, label_cols):\n    print(\"4. Loading IA Weights...\")\n    if not os.path.exists(ia_path):\n        print(\"   Warning: IA.tsv not found, using equal weights (1.0).\")\n        return torch.ones(len(label_cols), dtype=torch.float32)\n    \n    # IA.tsv không header, 2 cột: term và score\n    ia_df = pd.read_csv(ia_path, sep='\\t', header=None, names=['term', 'score'])\n    \n    weight_dict = dict(zip(ia_df['term'], ia_df['score']))\n    \n    weights = []\n    missing = 0\n    for col in label_cols:\n        w = weight_dict.get(col, 1.0)\n        weights.append(w)\n        if w == 1.0:\n            missing += 1\n    \n    if missing > 0:\n        print(f\"   {missing} labels not found in IA.tsv, using weight 1.0\")\n    \n    return torch.tensor(weights, dtype=torch.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.352128Z","iopub.execute_input":"2025-12-12T00:32:57.352634Z","iopub.status.idle":"2025-12-12T00:32:57.379977Z","shell.execute_reply.started":"2025-12-12T00:32:57.352589Z","shell.execute_reply":"2025-12-12T00:32:57.379139Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport os\nimport gc\n\nlabel_path = '/kaggle/working/label_p.parquet'\nif not os.path.exists(label_path):\n    print(\"Creating label_p.parquet with memory-efficient method (no crosstab)...\")\n\n    terms_path = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv'\n    \n    # Bước 1: Load chỉ 2 cột cần thiết\n    print(\"   Loading train_terms.tsv (only EntryID and term)...\")\n    terms = pd.read_csv(terms_path, sep='\\t', usecols=['EntryID', 'term'])\n    print(f\"   Loaded {terms.shape[0]} rows, memory ~{terms.memory_usage(deep=True).sum() / (1024**3):.2f} GB\")\n\n    # Bước 2: Lấy top_k terms phổ biến nhất\n    top_k = 1000  # Hoặc thay bằng cfg.top_k nếu đã có cfg\n    print(f\"   Selecting top {top_k} terms...\")\n    top_terms = terms['term'].value_counts().head(top_k).index\n    print(f\"   Top {top_k} terms selected.\")\n\n    # Bước 3: Chỉ giữ rows có term trong top_k (giảm dữ liệu mạnh)\n    print(\"   Filtering terms to top_k only...\")\n    terms_top = terms[terms['term'].isin(top_terms)].copy()\n    print(f\"   Filtered to {terms_top.shape[0]} rows\")\n    del terms\n    gc.collect()\n\n    # Bước 4: Dùng get_dummies (nhẹ hơn crosstab rất nhiều)\n    print(\"   Creating binary matrix with get_dummies...\")\n    label_matrix = pd.get_dummies(terms_top, columns=['term'], sparse=False)  # sparse=False để dễ xử lý\n    # Group by EntryID và max (để binary)\n    label_matrix = label_matrix.groupby('EntryID').max().astype('int8')\n    \n    # Reset index để có cột EntryID\n    label_matrix = label_matrix.reset_index()\n    \n    print(f\"   Final shape: {label_matrix.shape}, memory ~{label_matrix.memory_usage(deep=True).sum() / (1024**3):.2f} GB\")\n\n    # Bước 5: Save\n    label_matrix.to_parquet(label_path, index=False)\n    print(\"Done! label_p.parquet created successfully.\")\n\n    # Giải phóng RAM\n    del terms_top, label_matrix\n    gc.collect()\nelse:\n    print(\"label_p.parquet already exists.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.382625Z","iopub.execute_input":"2025-12-12T00:32:57.382915Z","iopub.status.idle":"2025-12-12T00:32:57.405027Z","shell.execute_reply.started":"2025-12-12T00:32:57.382896Z","shell.execute_reply":"2025-12-12T00:32:57.403938Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q biopython  # Chỉ install nếu cần\n\nimport pandas as pd\nimport os\nfrom Bio import SeqIO\nimport gc\n\nseq_path = '/kaggle/working/sequence.parquet'\nif not os.path.exists(seq_path):\n    print(\"Creating sequence.parquet safely (chunked to avoid OOM)...\")\n    \n    fasta_path = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\n    chunk_size = 10000  # Xử lý từng chunk nhỏ\n    chunks = []\n    \n    count = 0\n    current_chunk = []\n    \n    print(\"Parsing FASTA and saving in chunks...\")\n    for record in SeqIO.parse(fasta_path, 'fasta'):\n        current_chunk.append({'EntryID': record.id, 'seq': str(record.seq)})\n        count += 1\n        \n        if len(current_chunk) >= chunk_size:\n            chunks.append(pd.DataFrame(current_chunk))\n            print(f\"   Processed {count} sequences...\")\n            current_chunk = []  # Giải phóng chunk cũ\n            gc.collect()\n    \n    # Thêm chunk cuối\n    if current_chunk:\n        chunks.append(pd.DataFrame(current_chunk))\n        print(f\"   Processed {count} sequences (final chunk)...\")\n    \n    # Gộp tất cả chunk\n    print(\"Concatenating chunks...\")\n    df = pd.concat(chunks, ignore_index=True)\n    \n    # Chỉ giữ cột EntryID nếu không cần seq (tiết kiệm RAM lớn)\n    # Nếu code của bạn không dùng cột 'seq' (chỉ dùng EntryID), bỏ cột seq đi\n    if 'seq' in df.columns:\n        print(\"Dropping 'seq' column to save memory (only keep EntryID)...\")\n        df = df[['EntryID']]  # Code gốc chỉ cần EntryID để merge\n    \n    # Lưu\n    df.to_parquet(seq_path, index=False)\n    print(f\"sequence.parquet created successfully! Shape: {df.shape}\")\n    \n    # Giải phóng hoàn toàn\n    del chunks, current_chunk, df\n    gc.collect()\nelse:\n    print(\"sequence.parquet already exists.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:32:57.406121Z","iopub.execute_input":"2025-12-12T00:32:57.406412Z","iopub.status.idle":"2025-12-12T00:33:00.113038Z","shell.execute_reply.started":"2025-12-12T00:32:57.40639Z","shell.execute_reply":"2025-12-12T00:33:00.111613Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -q biopython\n\nimport pandas as pd\nimport os\nfrom Bio import SeqIO\nimport gc\n\nseq_path = '/kaggle/working/sequence.parquet'\nif not os.path.exists(seq_path):\n    print(\"Creating sequence.parquet with clean EntryID...\")\n    seqs = []\n    fasta_path = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\n    \n    for record in SeqIO.parse(fasta_path, 'fasta'):\n        # Extract only accession (e.g., 'A0A0C5B5G6' from 'sp|A0A0C5B5G6|MOTSC_HUMAN')\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) % 10000 == 0:\n            print(f\"Processed {len(seqs)} sequences...\")\n    \n    df = pd.DataFrame(seqs)\n    df.to_parquet(seq_path, index=False)\n    print(f\"sequence.parquet created! Shape: {df.shape}\")\n    print(\"Sample EntryID:\", df['EntryID'].head().tolist())\n    \n    del seqs, df\n    gc.collect()\nelse:\n    print(\"sequence.parquet exists. Checking and cleaning if needed...\")\n    df = pd.read_parquet(seq_path)\n    if df['EntryID'].iloc[0].startswith('sp|'):\n        print(\"Cleaning EntryID format...\")\n        df['EntryID'] = df['EntryID'].apply(lambda x: x.split('|')[1] if '|' in x else x)\n        df.to_parquet(seq_path, index=False)\n        print(\"Cleaned and resaved sequence.parquet.\")\n    print(\"Sample EntryID:\", df['EntryID'].head().tolist())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:33:00.114652Z","iopub.execute_input":"2025-12-12T00:33:00.115116Z","iopub.status.idle":"2025-12-12T00:33:02.98998Z","shell.execute_reply.started":"2025-12-12T00:33:00.115075Z","shell.execute_reply":"2025-12-12T00:33:02.988621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 8: Main Training Loop (sửa hoàn chỉnh - fix UnpicklingError + prediction ổn định)\n\n# Seed for reproducibility\nL.seed_everything(42)\n\n# 1. Load Train Data\nfull_df, embeddings, label_cols, tax_map, num_tax = load_and_align_train_data()\n\n# 2. Load Weights\nterm_weights = load_ia_weights(cfg.ia_path, label_cols)\n\n# 3. Split Train/Val\ntrain_idx, val_idx = train_test_split(np.arange(len(full_df)), test_size=0.2, random_state=42)\ntrain_df = full_df.iloc[train_idx].reset_index(drop=True)\nval_df = full_df.iloc[val_idx].reset_index(drop=True)\ntrain_emb = embeddings[train_idx]\nval_emb = embeddings[val_idx]\n\ndel full_df, embeddings\ngc.collect()\n\n# 4. Load Test Data\ntest_df, test_emb = load_test_data(tax_map)\n\n# 5. Setup Data Module\ndata_module = LitProteinDataModule(\n    train_df=train_df,\n    train_emb=train_emb,\n    val_df=val_df,\n    val_emb=val_emb,\n    test_df=test_df,\n    test_emb=test_emb,\n    label_columns=label_cols,\n    taxonomy_map=tax_map,\n    batch_size=cfg.batch_size,\n    num_workers=cfg.num_workers\n)\ndata_module.setup()\n\n# 6. Setup Model\nprint(f\"5. Init Model (SeqDim: {cfg.input_dim}, TaxDim: {cfg.taxonomy_dim}, Classes: {len(label_cols)})...\")\nmodel = LitMLPModule(\n    input_dim=cfg.input_dim,\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# 7. Trainer với MLFlowLogger (local mode)\ncheckpoint_callback = ModelCheckpoint(\n    dirpath=\"checkpoints\",\n    filename=\"best-multimodal-{epoch:02d}-{val_f1:.4f}\",\n    save_top_k=1,\n    monitor=\"val_f1\",\n    mode=\"max\",\n    save_weights_only=True\n)\n\nearly_stopping = EarlyStopping(monitor=\"val_f1\", patience=cfg.patience, mode=\"max\", verbose=True)\n\nmlf_logger = MLFlowLogger(\n    experiment_name=cfg.experiment_name,\n    tracking_uri=cfg.tracking_uri,  # None → lưu local tại /kaggle/working/mlruns\n    log_model=\"all\"\n)\n\ntrainer = L.Trainer(\n    max_epochs=cfg.num_epochs,\n    callbacks=[checkpoint_callback, early_stopping],\n    accelerator=\"auto\",\n    devices=\"auto\",\n    precision=\"16-mixed\" if torch.cuda.is_available() else 32,\n    logger=mlf_logger,\n    enable_progress_bar=True,\n    log_every_n_steps=10\n)\n\nprint(\"Start Training...\")\ntrainer.fit(model, data_module)\n\n# Log best model artifact\nbest_model_path = checkpoint_callback.best_model_path\nif best_model_path and os.path.exists(best_model_path):\n    print(f\"Best model saved at: {best_model_path}\")\n    print(\"Logging best model to MLflow...\")\n    mlf_logger.experiment.log_artifact(\n        run_id=mlf_logger.run_id,\n        local_path=best_model_path,\n        artifact_path=\"checkpoints\"\n    )\n    print(\"Best model logged as artifact.\")\nelse:\n    print(\"No best model checkpoint found.\")\n\n# 8. Prediction on Test Set\nif test_df is not None:\n    print(\"Predicting on Test Set...\")\n    \n    # FIX UnpicklingError: Cho phép SimpleNamespace khi load checkpoint\n    import torch\n    from types import SimpleNamespace\n    torch.serialization.add_safe_globals([SimpleNamespace])\n    \n    predictions_iterator = trainer.predict(\n        model, \n        dataloaders=data_module.test_dataloader(), \n        ckpt_path='best'\n    )\n    \n    print(f\"Writing predictions to {cfg.output}...\")\n    num_test_batches = len(data_module.test_dataloader())\n    \n    writer = None\n    try:\n        for batch in tqdm(predictions_iterator, total=num_test_batches, desc=\"Predicting\"):\n            preds_tensor = batch['predictions'].cpu().numpy()\n            ids = batch['ids']\n            \n            preds_df = pd.DataFrame(preds_tensor, columns=label_cols)\n            preds_df['EntryID'] = ids\n            \n            melted = preds_df.melt(id_vars='EntryID', var_name='term', value_name='score')\n            melted = melted[melted['score'] > 0.01]  # Threshold 0.01\n            \n            table = pa.Table.from_pandas(melted, preserve_index=False)\n            \n            if writer is None:\n                writer = pq.ParquetWriter(cfg.output, table.schema, compression='snappy')\n            writer.write_table(table)\n    finally:\n        if writer:\n            writer.close()\n    \n    print(f\"Done! Submission saved to {cfg.output}\")\nelse:\n    print(\"No test data loaded. Skipping prediction.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T00:33:02.991606Z","iopub.execute_input":"2025-12-12T00:33:02.991946Z","iopub.status.idle":"2025-12-12T00:49:03.491684Z","shell.execute_reply.started":"2025-12-12T00:33:02.99191Z","shell.execute_reply":"2025-12-12T00:49:03.490637Z"}},"outputs":[],"execution_count":null}]}