{"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":"!pip install -q --upgrade pip\n!pip install -q mlflow\n!pip install -q --no-cache-dir mlflow \n!pip install -q lightning torchmetrics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T16:16:48.650881Z","iopub.execute_input":"2025-12-12T16:16:48.651679Z","iopub.status.idle":"2025-12-12T16:16:57.073558Z","shell.execute_reply.started":"2025-12-12T16:16:48.651648Z","shell.execute_reply":"2025-12-12T16:16:57.072560Z"}},"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-12T16:16:57.075276Z","iopub.execute_input":"2025-12-12T16:16:57.075562Z","iopub.status.idle":"2025-12-12T16:16:57.080599Z","shell.execute_reply.started":"2025-12-12T16:16:57.075538Z","shell.execute_reply":"2025-12-12T16:16:57.079972Z"}},"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-12T16:16:57.081666Z","iopub.execute_input":"2025-12-12T16:16:57.082081Z","iopub.status.idle":"2025-12-12T16:16:57.098980Z","shell.execute_reply.started":"2025-12-12T16:16:57.082065Z","shell.execute_reply":"2025-12-12T16:16:57.098261Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 2: Config\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 \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-12T16:16:57.099804Z","iopub.execute_input":"2025-12-12T16:16:57.099980Z","iopub.status.idle":"2025-12-12T16:16:57.116609Z","shell.execute_reply.started":"2025-12-12T16:16:57.099966Z","shell.execute_reply":"2025-12-12T16:16:57.115948Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 4: Metrics\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-12T16:16:57.118488Z","iopub.execute_input":"2025-12-12T16:16:57.118682Z","iopub.status.idle":"2025-12-12T16:16:57.134494Z","shell.execute_reply.started":"2025-12-12T16:16:57.118666Z","shell.execute_reply":"2025-12-12T16:16:57.133728Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport numpy as np\nimport pandas as pd\nfrom torch.utils.data import Dataset, DataLoader\nimport lightning as L\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        # Tự động tìm cột Taxonomy (hỗ trợ cả TaxID và taxonomyID)\n        self.tax_col = None\n        if 'taxonomyID' in self.dataframe.columns:\n            self.tax_col = 'taxonomyID'\n        elif 'TaxID' in self.dataframe.columns:\n            self.tax_col = 'TaxID'\n            \n        # Default index cho unknown tax\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        protein_id = self.dataframe.iloc[index][self.id_col]\n        \n        # Đặt tên key là 'ids' để khớp với code predict\n        item = dict(\n            ids=protein_id, \n            embedding=torch.tensor(emb_vector, dtype=torch.float32)\n        )\n        \n        # 2. Taxonomy Indexing\n        tax_idx = self.unknown_tax_idx\n        if self.taxonomy_map is not None and self.tax_col is not None:\n            raw_tax = self.dataframe.iloc[index][self.tax_col]\n            tax_idx = self.taxonomy_map.get(raw_tax, self.unknown_tax_idx)\n            \n        # Đặt tên key là 'taxonomy' để khớp với Model forward\n        item['taxonomy'] = torch.tensor(tax_idx, dtype=torch.long)\n        \n        # 3. Labels\n        if self.label_columns is not None:\n            labels = self.dataframe.iloc[index][self.label_columns].values.astype(float)\n            item['labels'] = torch.tensor(labels, dtype=torch.float32)\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-12T16:16:57.135324Z","iopub.execute_input":"2025-12-12T16:16:57.135562Z","iopub.status.idle":"2025-12-12T16:16:57.154746Z","shell.execute_reply.started":"2025-12-12T16:16:57.135547Z","shell.execute_reply":"2025-12-12T16:16:57.154057Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport lightning as L\nfrom torchmetrics.classification import MultilabelF1Score\n\n# 1. Residual Block (Giữ nguyên)\nclass ResidualBlock(nn.Module):\n    def __init__(self, hidden_dim, dropout_rate=0.2):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout_rate),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim)\n        )\n        self.relu = nn.ReLU()\n\n    def forward(self, x):\n        residual = x\n        out = self.block(x)\n        out += residual\n        return self.relu(out)\n\n# 2. CẬP NHẬT FOCAL LOSS (Dùng Logits để an toàn Mixed Precision)\nclass FocalLoss(nn.Module):\n    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):\n        super().__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.reduction = reduction\n\n    def forward(self, logits, targets):\n        # binary_cross_entropy_with_logits an toàn cho 16-mixed precision\n        bce_loss = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')\n        pt = torch.exp(-bce_loss)\n        f_loss = self.alpha * (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\n# 3. CẬP NHẬT MODEL (Trả về Logits ở forward)\nclass LitMLPModule(L.LightningModule):\n    def __init__(self, input_dim, taxonomy_num, taxonomy_dim, hidden_dim, num_labels, learning_rate, term_weights, cfg):\n        super().__init__()\n        self.save_hyperparameters()\n        self.cfg = cfg\n        \n        self.protein_encoder = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU()\n        )\n        \n        self.taxonomy_embedding = nn.Embedding(taxonomy_num, taxonomy_dim)\n        self.taxonomy_encoder = nn.Sequential(\n            nn.Linear(taxonomy_dim, taxonomy_dim),\n            nn.ReLU()\n        )\n        \n        combined_dim = hidden_dim + taxonomy_dim\n        self.projection = nn.Linear(combined_dim, hidden_dim)\n        \n        self.res_blocks = nn.Sequential(\n            ResidualBlock(hidden_dim, dropout_rate=0.3),\n            ResidualBlock(hidden_dim, dropout_rate=0.3),\n            ResidualBlock(hidden_dim, dropout_rate=0.3)\n        )\n        \n        self.classifier = nn.Linear(hidden_dim, num_labels)\n        \n        self.loss_fn = FocalLoss(alpha=0.25, gamma=2.0)\n        self.val_f1 = MultilabelF1Score(num_labels=num_labels, threshold=0.3, average='macro')\n\n    def forward(self, x, tax_indices):\n        prot_emb = self.protein_encoder(x)\n        tax_emb = self.taxonomy_embedding(tax_indices)\n        tax_emb = self.taxonomy_encoder(tax_emb)\n        \n        combined = torch.cat([prot_emb, tax_emb], dim=1)\n        x = self.projection(combined)\n        x = self.res_blocks(x)\n        \n        logits = self.classifier(x)\n        # SỬA QUAN TRỌNG: Không dùng sigmoid ở đây nữa, trả về logits thô\n        return logits \n\n    def training_step(self, batch, batch_idx):\n        x, y, tax = batch['embedding'], batch['labels'], batch['taxonomy']\n        logits = self(x, tax)\n        loss = self.loss_fn(logits, y.float())\n        self.log(\"train_loss\", loss, prog_bar=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y, tax = batch['embedding'], batch['labels'], batch['taxonomy']\n        logits = self(x, tax)\n        loss = self.loss_fn(logits, y.float())\n        \n        # Khi tính F1 thì mới dùng Sigmoid để ra xác suất\n        preds = torch.sigmoid(logits)\n        self.val_f1(preds, y)\n        self.log(\"val_loss\", loss, prog_bar=True)\n        self.log(\"val_f1\", self.val_f1, prog_bar=True)\n\n    def predict_step(self, batch, batch_idx):\n        x, tax = batch['embedding'], batch['taxonomy']\n        ids = batch['ids']\n        logits = self(x, tax)\n        # Khi dự đoán cuối cùng cũng cần Sigmoid\n        preds = torch.sigmoid(logits)\n        return {'ids': ids, 'predictions': preds}\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.learning_rate, weight_decay=1e-4)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer, mode='max', factor=0.5, patience=3, verbose=True\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"monitor\": \"val_f1\"\n            }\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T16:16:57.155404Z","iopub.execute_input":"2025-12-12T16:16:57.155637Z","iopub.status.idle":"2025-12-12T16:16:57.173819Z","shell.execute_reply.started":"2025-12-12T16:16:57.155622Z","shell.execute_reply":"2025-12-12T16:16:57.173145Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 7: Train Functions\n\ndef load_and_align_train_data():\n    print(\"1. Loading Training Data...\")\n    \n    # A. Load Sequence\n    print(f\"   Loading Sequence from {cfg.sequence_path}...\")\n    seq_df = pd.read_parquet(cfg.sequence_path)\n    \n    # B. Load Taxonomy Train \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    # Normalize EntryID \n    seq_df['EntryID'] = seq_df['EntryID'].apply(lambda x: x.split('|')[1] if '|' in str(x) else x)\n    lbl_df['EntryID'] = lbl_df['EntryID'].apply(lambda x: x.split('|')[1] if '|' in str(x) else x)\n    # Inner merge để chỉ giữ protein có nhãn\n    full_df = seq_df.merge(lbl_df, on='EntryID', how='inner')\n    if len(full_df) == 0:\n        raise ValueError(\"No overlapping EntryID between sequence and labels. Check formats!\")\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\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        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 \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-12T16:16:57.174616Z","iopub.execute_input":"2025-12-12T16:16:57.174844Z","iopub.status.idle":"2025-12-12T16:16:57.194761Z","shell.execute_reply.started":"2025-12-12T16:16:57.174829Z","shell.execute_reply":"2025-12-12T16:16:57.194075Z"}},"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    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    top_k = 1000 \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    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    print(\"   Creating binary matrix with get_dummies...\")\n    label_matrix = pd.get_dummies(terms_top, columns=['term'], sparse=False) \n    \n    label_matrix = label_matrix.groupby('EntryID').max().astype('int8')\n    \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    label_matrix.to_parquet(label_path, index=False)\n    print(\"Done! label_p.parquet created successfully.\")\n\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-12T16:16:57.195539Z","iopub.execute_input":"2025-12-12T16:16:57.195733Z","iopub.status.idle":"2025-12-12T16:16:57.211819Z","shell.execute_reply.started":"2025-12-12T16:16:57.195711Z","shell.execute_reply":"2025-12-12T16:16:57.211251Z"}},"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 safely (chunked to avoid OOM)...\")\n    \n    fasta_path = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\n    chunk_size = 10000 \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 = [] \n            gc.collect()\n    \n    if current_chunk:\n        chunks.append(pd.DataFrame(current_chunk))\n        print(f\"   Processed {count} sequences (final chunk)...\")\n    \n    print(\"Concatenating chunks...\")\n    df = pd.concat(chunks, ignore_index=True)\n    \n    if 'seq' in df.columns:\n        print(\"Dropping 'seq' column to save memory (only keep EntryID)...\")\n        df = df[['EntryID']]  \n    \n    df.to_parquet(seq_path, index=False)\n    print(f\"sequence.parquet created successfully! Shape: {df.shape}\")\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-12T16:16:57.212555Z","iopub.execute_input":"2025-12-12T16:16:57.212853Z","iopub.status.idle":"2025-12-12T16:16:59.185398Z","shell.execute_reply.started":"2025-12-12T16:16:57.212831Z","shell.execute_reply":"2025-12-12T16:16:59.184404Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\n# Load seq_df và lbl_df\nseq_df = pd.read_parquet('/kaggle/working/sequence.parquet')\nlbl_df = pd.read_parquet('/kaggle/working/label_p.parquet')\n\nprint(f\"seq_df shape: {seq_df.shape}, columns: {seq_df.columns.tolist()}\")\nprint(f\"lbl_df shape: {lbl_df.shape}, columns: {lbl_df.columns.tolist()}\")\n\n# Kiểm tra overlap EntryID\ncommon = set(seq_df['EntryID']).intersection(set(lbl_df['EntryID']))\nprint(f\"Common EntryID: {len(common)}\")\n\nif len(common) == 0:\n    print(\"NO OVERLAP! Sample EntryID seq_df:\", seq_df['EntryID'].head().tolist())\n    print(\"Sample EntryID lbl_df:\", lbl_df['EntryID'].head().tolist())\nelse:\n    print(\"Overlap OK! Common samples:\", common)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T16:16:59.186593Z","iopub.execute_input":"2025-12-12T16:16:59.186865Z","iopub.status.idle":"2025-12-12T16:16:59.446437Z","shell.execute_reply.started":"2025-12-12T16:16:59.186842Z","shell.execute_reply":"2025-12-12T16:16:59.445501Z"}},"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\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-12T16:16:59.447355Z","iopub.execute_input":"2025-12-12T16:16:59.447581Z","iopub.status.idle":"2025-12-12T16:17:01.451349Z","shell.execute_reply.started":"2025-12-12T16:16:59.447561Z","shell.execute_reply":"2025-12-12T16:17:01.450594Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Cell 8: Main Training Loop\n\nimport os\nimport pandas as pd\nimport torch\nfrom pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping\nfrom pytorch_lightning.loggers import MLFlowLogger\nimport lightning as L\nimport gc\nimport numpy as np\nfrom sklearn.model_selection import train_test_split\nfrom types import SimpleNamespace\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\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,\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=False, \n    log_every_n_steps=50       \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}\")\nelse:\n    print(\"No best model checkpoint found.\")\n\nif test_df is not None:\n    print(\"Predicting on Test Set...\")\n    \n    # FIX UnpicklingError\n    torch.serialization.add_safe_globals([SimpleNamespace])\n    \n    # Lấy DataLoader\n    test_loader = data_module.test_dataloader()\n    \n    # Sử dụng trainer.predict\n    predictions_iterator = trainer.predict(\n        model, \n        dataloaders=test_loader, \n        ckpt_path='best'\n    )\n    \n    output_tsv_path = 'submission.tsv'\n    print(f\"Writing predictions directly to {output_tsv_path}...\")\n    \n    # --- PHẦN SỬA CHÍNH: Xử lý file output ---\n    # 1. Tạo file rỗng/Xóa nội dung cũ\n    with open(output_tsv_path, 'w') as f:\n        pass \n\n    try:\n        num_test_batches = len(test_loader)\n    except:\n        num_test_batches = \"Unknown\"\n    \n    for i, batch in enumerate(predictions_iterator):\n        if i % 100 == 0:\n            print(f\"Processing batch {i}/{num_test_batches}...\", flush=True)\n\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        # Melt dữ liệu\n        melted = preds_df.melt(id_vars='EntryID', var_name='term', value_name='score')\n        \n        # --- SỬA LOGIC Ở ĐÂY (CLEAN PREFIX & ROUNDING) ---\n        \n        # 1. Xóa chữ 'term_' nếu có trong cột 'term'\n        melted['term'] = melted['term'].str.replace('term_', '', regex=False)\n        \n        # 2. Làm tròn score lấy 3 chữ số thập phân\n        melted['score'] = melted['score'].round(3)\n        \n        # 3. Lọc score > 0.01 (Giữ lại các dự đoán có độ tin cậy nhất định)\n        melted = melted[melted['score'] > 0.01]\n        \n        if not melted.empty:\n            # 4. Luôn append và KHÔNG ghi header\n            melted.to_csv(\n                output_tsv_path, \n                sep='\\t', \n                index=False, \n                mode='a',     # Luôn nối tiếp\n                header=False  # KHÔNG BAO GIỜ GHI HEADER\n            )\n            \n    print(f\"Done! Submission saved to {output_tsv_path}\")\n    \n    # --- IN 5 DÒNG ĐẦU CỦA SUBMISSION ---\n    print(\"Preview (First 5 lines):\")\n    # Đọc với header=None để kiểm tra dòng đầu tiên có phải là dữ liệu không\n    print(pd.read_csv(output_tsv_path, sep='\\t', header=None).head())\n\nelse:\n    print(\"No test data loaded. Skipping prediction.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T16:17:01.452452Z","iopub.execute_input":"2025-12-12T16:17:01.452703Z","iopub.status.idle":"2025-12-12T16:41:10.449715Z","shell.execute_reply.started":"2025-12-12T16:17:01.452678Z","shell.execute_reply":"2025-12-12T16:41:10.448982Z"}},"outputs":[],"execution_count":null}]}