{"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":14084779,"sourceType":"competition"},{"sourceId":5549164,"sourceType":"datasetVersion","datasetId":3197305},{"sourceId":5792099,"sourceType":"datasetVersion","datasetId":3327296},{"sourceId":6247561,"sourceType":"datasetVersion","datasetId":3590060}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport time\nimport matplotlib.pyplot as plt\nimport os\n\n# TORCH MODULES\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torchmetrics.classification import MultilabelF1Score","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# --- 1. SETUP & CONFIGURATION ---","metadata":{}},{"cell_type":"code","source":"# All constants and hyperparameters are defined in one place for easy management.\nclass config:\n    MAIN_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\n    \n    # Model and Training Parameters\n    num_labels = 500\n    n_epochs = 8\n    batch_size = 128\n    lr = 0.001  # Lowered learning rate for a more stable start.\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f\"Using device: {config.device}\")\n\n# Mapping for embeddings\nembeds_map = {\n    \"T5\" : \"t5embeds\",\n    \"ProtBERT\" : \"protbert-embeddings-for-cafa5\",\n    \"EMS2\" : \"cafa-5-ems-2-embeddings-numpy\"\n}\nembeds_dim = {\n    \"T5\" : 1024,\n    \"ProtBERT\" : 1024,\n    \"EMS2\" : 1280\n}","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# --- 2. DATA LOADING: PROTEIN SEQUENCE DATASET ---","metadata":{}},{"cell_type":"code","source":"# This PyTorch Dataset class is responsible for loading the protein embeddings and \n# their corresponding labels. It handles both training data (embeddings + labels) \n# and test data (embeddings + protein IDs).\nclass ProteinSequenceDataset(Dataset):\n    def __init__(self, datatype, embeddings_source):\n        super(ProteinSequenceDataset).__init__()\n        self.datatype = datatype\n        \n        base_path = f\"/kaggle/input/{embeds_map[embeddings_source]}/\"\n        \n        # Construct paths correctly\n        embeds_path = os.path.join(base_path, f\"{datatype}_embeddings.npy\")\n        ids_path = os.path.join(base_path, f\"{datatype}_ids.npy\")\n        \n        # Handle special filenames for T5\n        if embeddings_source == \"T5\":\n            embeds_path = os.path.join(base_path, f\"{datatype}_embeds.npy\")\n\n        embeds = np.load(embeds_path)\n        ids = np.load(ids_path)\n            \n        # Optimized DataFrame creation\n        self.df = pd.DataFrame({\"EntryID\": ids, \"embed\": list(embeds)})\n        \n        if datatype==\"train\":\n            labels_path = f\"/kaggle/input/train-targets-top{config.num_labels}/train_targets_top{config.num_labels}.npy\"\n            np_labels = np.load(labels_path)\n            \n            # Merge labels into the DataFrame\n            labels_df = pd.DataFrame({\"EntryID\": self.df['EntryID'], \"labels_vect\": list(np_labels)})\n            self.df = pd.merge(self.df, labels_df, on=\"EntryID\")\n            \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, index):\n        embed = torch.tensor(self.df.iloc[index][\"embed\"], dtype=torch.float32)\n        if self.datatype == \"train\":\n            targets = torch.tensor(self.df.iloc[index][\"labels_vect\"], dtype=torch.float32)\n            return embed, targets\n        else: # datatype == \"test\"\n            protein_id = self.df.iloc[index][\"EntryID\"]\n            return embed, protein_id","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# --- 3. MODEL ARCHITECTURE: 1D CONVOLUTIONAL NEURAL NETWORK (CNN) ---","metadata":{}},{"cell_type":"code","source":"# The model is a 1D CNN designed to find patterns in the sequence embeddings. Key features include:\n# - `Conv1d` layers for feature extraction.\n# - `BatchNorm1d` to stabilize and speed up training.\n# - `ReLU` as the activation function.\n# - `MaxPool1d` for down-sampling.\n# - `Dropout` in the final fully-connected block to prevent overfitting.\nclass CNN1D(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super(CNN1D, self).__init__()\n        self.conv_block1 = nn.Sequential(\n            nn.Conv1d(in_channels=1, out_channels=32, kernel_size=5, padding=2),\n            nn.BatchNorm1d(32),\n            nn.ReLU(),\n            nn.MaxPool1d(kernel_size=2, stride=2)\n        )\n        self.conv_block2 = nn.Sequential(\n            nn.Conv1d(in_channels=32, out_channels=64, kernel_size=3, padding=1),\n            nn.BatchNorm1d(64),\n            nn.ReLU(),\n            nn.MaxPool1d(kernel_size=2, stride=2)\n        )\n        \n        flattened_size = int(64 * input_dim / 4)\n        \n        self.fc_block = nn.Sequential(\n            nn.Linear(in_features=flattened_size, out_features=1024),\n            nn.ReLU(),\n            nn.Dropout(p=0.4), # Added Dropout to prevent overfitting\n            nn.Linear(in_features=1024, out_features=num_classes)\n        )\n\n    def forward(self, x):\n        # (batch_size, embed_size) -> (batch_size, 1, embed_size)\n        x = x.unsqueeze(1) \n        x = self.conv_block1(x)\n        x = self.conv_block2(x)\n        x = torch.flatten(x, 1)\n        x = self.fc_block(x)\n        return x","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# --- 4. MODEL TRAINING ---","metadata":{}},{"cell_type":"code","source":"# This function encapsulates the entire training and validation loop. Key improvements include:\n# - Loss Function: Using `BCEWithLogitsLoss`, which is essential for multi-label classification.\n# - Metric-driven Learning: The learning rate scheduler adjusts based on the validation F1-score.\n# - Optimal Threshold Finding: During validation, the code iterates through different thresholds \n#   to find the one that maximizes the F1-score for the current epoch.\ndef train_model(embeddings_source, model_type=\"convolutional\", train_size=0.9):\n    \n    train_dataset = ProteinSequenceDataset(datatype=\"train\", embeddings_source=embeddings_source)\n    \n    train_set, val_set = random_split(train_dataset, \n                                      lengths=[int(len(train_dataset) * train_size), \n                                               len(train_dataset) - int(len(train_dataset) * train_size)])\n    \n    train_dataloader = DataLoader(train_set, batch_size=config.batch_size, shuffle=True)\n    val_dataloader = DataLoader(val_set, batch_size=config.batch_size, shuffle=False)\n\n    if model_type == \"convolutional\":\n        model = CNN1D(input_dim=embeds_dim[embeddings_source], num_classes=config.num_labels).to(config.device)\n    else:\n        raise ValueError(\"Unsupported model type\")\n\n    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)\n    scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.1, patience=1, verbose=True) # Operates on F1-score\n    \n    # CRITICAL CHANGE: Using the correct loss function for multi-label classification\n    loss_fn = torch.nn.BCEWithLogitsLoss()\n    \n    # Metrics\n    f1_metric = MultilabelF1Score(num_labels=config.num_labels, average='macro').to(config.device)\n\n    print(\"STARTING TRAINING...\")\n    \n    best_val_f1 = 0.0\n    best_threshold = 0.5\n\n    for epoch in range(config.n_epochs):\n        print(f\"EPOCH {epoch+1}/{config.n_epochs}\")\n        \n        # --- TRAIN PHASE ---\n        model.train()\n        total_train_loss = 0\n        for embed, targets in tqdm(train_dataloader, desc=\"Training\"):\n            embed, targets = embed.to(config.device), targets.to(config.device)\n            \n            optimizer.zero_grad()\n            preds_logits = model(embed)\n            loss = loss_fn(preds_logits, targets)\n            \n            loss.backward()\n            optimizer.step()\n            total_train_loss += loss.item()\n\n        avg_train_loss = total_train_loss / len(train_dataloader)\n        print(f\"Average Training Loss: {avg_train_loss:.4f}\")\n\n        # --- VALIDATION PHASE ---\n        model.eval()\n        all_val_preds = []\n        all_val_targets = []\n        with torch.no_grad():\n            for embed, targets in tqdm(val_dataloader, desc=\"Validation\"):\n                embed, targets = embed.to(config.device), targets.to(config.device)\n                preds_logits = model(embed)\n                \n                all_val_preds.append(torch.sigmoid(preds_logits))\n                all_val_targets.append(targets)\n\n        all_val_preds = torch.cat(all_val_preds)\n        all_val_targets = torch.cat(all_val_targets)\n\n        # Find the best F1 score by searching for the optimal threshold\n        best_f1_for_epoch = 0\n        best_thresh_for_epoch = 0\n        thresholds = np.arange(0.1, 0.51, 0.05)\n        for thresh in thresholds:\n            f1_metric.threshold = thresh\n            f1 = f1_metric(all_val_preds, all_val_targets.int())\n            if f1 > best_f1_for_epoch:\n                best_f1_for_epoch = f1\n                best_thresh_for_epoch = thresh\n\n        print(f\"Average Validation F1-Score: {best_f1_for_epoch:.4f} (at best threshold: {best_thresh_for_epoch:.2f})\")\n\n        scheduler.step(best_f1_for_epoch)\n\n        # Save the best model\n        if best_f1_for_epoch > best_val_f1:\n            best_val_f1 = best_f1_for_epoch\n            best_threshold = best_thresh_for_epoch\n            torch.save(model.state_dict(), \"best_model.pth\")\n            print(f\"New best model saved! F1: {best_val_f1:.4f}\")\n\n    print(\"\\nTRAINING FINISHED\")\n    print(f\"Highest Validation F1-Score: {best_val_f1:.4f}\")\n    print(f\"Best threshold for this score: {best_threshold:.2f}\")\n\n    # Load the best performing model\n    model.load_state_dict(torch.load(\"best_model.pth\"))\n    \n    return model, best_threshold\n\n# Train the model\nems2_model, best_threshold = train_model(embeddings_source=\"EMS2\", model_type=\"convolutional\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# --- 5. GENERATING PREDICTIONS ---","metadata":{}},{"cell_type":"code","source":"# After training, the best model is used to generate predictions on the test set. \n# The saved optimal threshold from the validation phase is used to convert model \n# probabilities into final binary predictions.\ndef predict(model, embeddings_source, threshold):\n    test_dataset = ProteinSequenceDataset(datatype=\"test\", embeddings_source=embeddings_source)\n    test_dataloader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)\n    \n    model.eval()\n    \n    # Pre-load label names\n    labels_df = pd.read_csv(os.path.join(config.MAIN_DIR, \"Train/train_terms.tsv\"), sep=\"\\t\")\n    top_terms = labels_df.groupby(\"term\")[\"EntryID\"].count().sort_values(ascending=False)\n    labels_names = top_terms.head(config.num_labels).index.values\n    \n    print(\"\\nGENERATING PREDICTIONS FOR THE TEST SET...\")\n    \n    results = []\n    with torch.no_grad():\n        for embed, ids in tqdm(test_dataloader, desc=\"Predicting\"):\n            embed = embed.to(config.device)\n            preds_logits = model(embed)\n            preds_probs = torch.sigmoid(preds_logits).cpu().numpy()\n            \n            # Collect predictions that are above the optimal threshold\n            for i, protein_id in enumerate(ids):\n                protein_probs = preds_probs[i]\n                go_indices = np.where(protein_probs > threshold)[0]\n                for idx in go_indices:\n                    results.append({\n                        \"Id\": protein_id,\n                        \"GO term\": labels_names[idx],\n                        \"Confidence\": protein_probs[idx]\n                    })\n    \n    submission_df = pd.DataFrame(results)\n    print(\"PREDICTIONS COMPLETE.\")\n    return submission_df\n\nsubmission_df = predict(ems2_model, \"EMS2\", best_threshold)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# --- 6. SUBMISSION FILE GENERATION ---","metadata":{}},{"cell_type":"code","source":"# The final step is to create the `submission.tsv` file. This involves merging our\n# model's predictions with an external, pre-existing submission file. The `.fillna()`\n# method provides an efficient way to combine them.\nprint(\"\\nMerging submission files...\")\n\n# Load external submission file\nsubmission2 = pd.read_csv('/kaggle/input/blast-quick-sprof-zero-pred/submission.tsv',\n                          sep='\\t', header=None, names=['Id', 'GO term', 'Confidence2'])\n\n# Merge the two submissions. `outer` join keeps all rows from both DataFrames.\nsubs = pd.merge(submission_df, submission2, on=['Id', 'GO term'], how='outer')\n\n# Combine confidence scores efficiently.\n# Fill NaN values in 'Confidence2' with values from our model's 'Confidence'.\nsubs['Confidence_combined'] = subs['Confidence2'].fillna(subs['Confidence'])\n\n# Select only the required columns and save to file\nfinal_submission = subs[['Id', 'GO term', 'Confidence_combined']]\nfinal_submission.to_csv('submission.tsv', sep='\\t', header=False, index=False)\n\nprint(\"Submission file 'submission.tsv' created successfully!\")\nprint(f\"It contains {len(final_submission)} predictions in total.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}