{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"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":30527,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n\"\"\"\nCAFA 6 Protein Function Prediction - Complete Implementation\nThis code predicts Gene Ontology (GO) terms for proteins based on their sequences\n\"\"\"\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport time\nimport os\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# PyTorch modules\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torchmetrics.classification import MultilabelF1Score\n\n# Visualization\nimport matplotlib.pyplot as plt\nplt.style.use('ggplot')\n\n# Set random seeds for reproducibility\nnp.random.seed(42)\ntorch.manual_seed(42)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(42)\n\n# ====================== CONFIGURATION ======================\nclass Config:\n    \"\"\"Configuration settings for the model\"\"\"\n    def __init__(self):\n        self.main_dir = \"/kaggle/input/cafa-6-protein-function-prediction\"\n        self.train_sequences_path = f\"{self.main_dir}/Train/train_sequences.fasta\"\n        self.train_labels_path = f\"{self.main_dir}/Train/train_terms.tsv\"\n        self.test_sequences_path = f\"{self.main_dir}/Test/testsuperset.fasta\"\n        self.ia_path = f\"{self.main_dir}/IA.tsv\"\n        \n        # Model parameters\n        self.num_labels = 600  # Start with top 500, can be increased\n        self.n_epochs = 15\n        self.batch_size = 64\n        self.lr = 0.001\n        self.dropout_rate = 0.3\n        self.patience = 3\n        \n        # Device configuration\n        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n        \n        # Embedding configurations\n        self.embeds_map = {\n            \"ESM2\": \"cafa-5-ems-2-embeddings-numpy\",\n            \"ProtBERT\": \"protbert-embeddings-for-cafa5\",\n            \"T5\": \"t5embeds\"\n        }\n        \n        self.embeds_dim = {\n            \"ESM2\": 1280,\n            \"ProtBERT\": 1024,\n            \"T5\": 1024\n        }\n\nconfig = Config()\nprint(f\"Using device: {config.device}\")\n\n# ====================== DATA LOADING ======================\nclass ProteinDataset(Dataset):\n    \"\"\"Dataset class for protein sequences and their GO annotations\"\"\"\n    \n    def __init__(self, datatype, embeddings_source, config):\n        super().__init__()\n        self.datatype = datatype\n        self.config = config\n        self.embeddings_source = embeddings_source\n        \n        # Load embeddings\n        self._load_embeddings()\n        \n        # Load labels for training data\n        if self.datatype == \"train\":\n            self._load_labels()\n    \n    def _load_embeddings(self):\n        \"\"\"Load pre-computed embeddings\"\"\"\n        embed_dir = f\"/kaggle/input/{self.config.embeds_map[self.embeddings_source]}\"\n        \n        if self.embeddings_source == \"ESM2\":\n            embeds = np.load(f\"{embed_dir}/{self.datatype}_embeddings.npy\")\n            ids = np.load(f\"{embed_dir}/{self.datatype}_ids.npy\")\n        elif self.embeddings_source == \"ProtBERT\":\n            embeds = np.load(f\"{embed_dir}/{self.datatype}_embeddings.npy\")\n            ids = np.load(f\"{embed_dir}/{self.datatype}_ids.npy\")\n        elif self.embeddings_source == \"T5\":\n            embeds = np.load(f\"{embed_dir}/{self.datatype}_embeds.npy\")\n            ids = np.load(f\"{embed_dir}/{self.datatype}_ids.npy\")\n        \n        # Create DataFrame\n        self.df = pd.DataFrame({\n            \"EntryID\": ids,\n            \"embed\": [embeds[i] for i in range(embeds.shape[0])]\n        })\n    \n    def _load_labels(self):\n        \"\"\"Load GO term labels for training data\"\"\"\n        # Load pre-processed top labels if available\n        label_file = f\"/kaggle/input/train-targets-top{self.config.num_labels}/train_targets_top{self.config.num_labels}.npy\"\n        \n        if os.path.exists(label_file):\n            np_labels = np.load(label_file)\n            df_labels = pd.DataFrame({\n                'EntryID': self.df['EntryID'],\n                'labels_vect': [row for row in np_labels]\n            })\n            self.df = self.df.merge(df_labels, on=\"EntryID\", how=\"inner\")\n        else:\n            # Process labels from scratch\n            self._process_labels_from_tsv()\n    \n    def _process_labels_from_tsv(self):\n        \"\"\"Process labels from the TSV file\"\"\"\n        labels_df = pd.read_csv(self.config.train_labels_path, sep=\"\\t\", names=[\"EntryID\", \"term\", \"aspect\"])\n        \n        # Get top terms\n        top_terms = labels_df.groupby(\"term\")[\"EntryID\"].count().sort_values(ascending=False)\n        self.top_terms = top_terms[:self.config.num_labels].index.tolist()\n        \n        # Create label vectors\n        label_vectors = []\n        for entry_id in self.df['EntryID']:\n            entry_terms = labels_df[labels_df['EntryID'] == entry_id]['term'].tolist()\n            vector = [1 if term in entry_terms else 0 for term in self.top_terms]\n            label_vectors.append(vector)\n        \n        self.df['labels_vect'] = label_vectors\n    \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        embed = torch.tensor(row[\"embed\"], dtype=torch.float32)\n        \n        if self.datatype == \"train\":\n            labels = torch.tensor(row[\"labels_vect\"], dtype=torch.float32)\n            return embed, labels\n        else:\n            return embed, row[\"EntryID\"]\n\n# ====================== MODEL ARCHITECTURES ======================\nclass ImprovedMLP(nn.Module):\n    \"\"\"Improved Multi-Layer Perceptron with dropout and batch norm\"\"\"\n    \n    def __init__(self, input_dim, num_classes, dropout_rate=0.3):\n        super().__init__()\n        \n        self.fc1 = nn.Linear(input_dim, 1024)\n        self.bn1 = nn.BatchNorm1d(1024)\n        self.dropout1 = nn.Dropout(dropout_rate)\n        \n        self.fc2 = nn.Linear(1024, 512)\n        self.bn2 = nn.BatchNorm1d(512)\n        self.dropout2 = nn.Dropout(dropout_rate)\n        \n        self.fc3 = nn.Linear(512, 256)\n        self.bn3 = nn.BatchNorm1d(256)\n        self.dropout3 = nn.Dropout(dropout_rate)\n        \n        self.fc4 = nn.Linear(256, num_classes)\n    \n    def forward(self, x):\n        x = F.relu(self.bn1(self.fc1(x)))\n        x = self.dropout1(x)\n        \n        x = F.relu(self.bn2(self.fc2(x)))\n        x = self.dropout2(x)\n        \n        x = F.relu(self.bn3(self.fc3(x)))\n        x = self.dropout3(x)\n        \n        x = self.fc4(x)\n        return x\n\nclass AttentionCNN(nn.Module):\n    \"\"\"1D CNN with attention mechanism\"\"\"\n    \n    def __init__(self, input_dim, num_classes, dropout_rate=0.3):\n        super().__init__()\n        \n        # Convolutional layers\n        self.conv1 = nn.Conv1d(1, 64, kernel_size=7, padding=3)\n        self.bn1 = nn.BatchNorm1d(64)\n        self.pool1 = nn.MaxPool1d(2)\n        \n        self.conv2 = nn.Conv1d(64, 128, kernel_size=5, padding=2)\n        self.bn2 = nn.BatchNorm1d(128)\n        self.pool2 = nn.MaxPool1d(2)\n        \n        self.conv3 = nn.Conv1d(128, 256, kernel_size=3, padding=1)\n        self.bn3 = nn.BatchNorm1d(256)\n        self.pool3 = nn.AdaptiveAvgPool1d(1)\n        \n        # Attention layer\n        self.attention = nn.MultiheadAttention(256, num_heads=8, dropout=dropout_rate, batch_first=True)\n        \n        # Fully connected layers\n        self.fc1 = nn.Linear(256, 512)\n        self.dropout = nn.Dropout(dropout_rate)\n        self.fc2 = nn.Linear(512, num_classes)\n    \n    def forward(self, x):\n        # Reshape for Conv1d: (batch, 1, features)\n        x = x.unsqueeze(1)\n        \n        # Convolutional blocks\n        x = F.relu(self.bn1(self.conv1(x)))\n        x = self.pool1(x)\n        \n        x = F.relu(self.bn2(self.conv2(x)))\n        x = self.pool2(x)\n        \n        x = F.relu(self.bn3(self.conv3(x)))\n        \n        # Reshape for attention: (batch, seq_len, features)\n        x = x.transpose(1, 2)\n        \n        # Self-attention\n        x, _ = self.attention(x, x, x)\n        \n        # Global pooling\n        x = x.mean(dim=1)\n        \n        # Classification head\n        x = F.relu(self.fc1(x))\n        x = self.dropout(x)\n        x = self.fc2(x)\n        \n        return x\n\nclass HybridModel(nn.Module):\n    \"\"\"Hybrid model combining CNN and LSTM features\"\"\"\n    \n    def __init__(self, input_dim, num_classes, dropout_rate=0.3):\n        super().__init__()\n        \n        # Feature extraction branch 1: CNN\n        self.conv1 = nn.Conv1d(1, 32, kernel_size=5, padding=2)\n        self.conv2 = nn.Conv1d(32, 64, kernel_size=3, padding=1)\n        self.pool = nn.AdaptiveMaxPool1d(128)\n        \n        # Feature extraction branch 2: LSTM\n        self.lstm = nn.LSTM(input_dim, 256, num_layers=2, \n                            bidirectional=True, dropout=dropout_rate, batch_first=True)\n        \n        # Fusion and classification\n        self.fc1 = nn.Linear(64 * 128 + 512, 512)\n        self.dropout = nn.Dropout(dropout_rate)\n        self.fc2 = nn.Linear(512, 256)\n        self.fc3 = nn.Linear(256, num_classes)\n    \n    def forward(self, x):\n        batch_size = x.shape[0]\n        \n        # CNN branch\n        cnn_x = x.unsqueeze(1)\n        cnn_x = F.relu(self.conv1(cnn_x))\n        cnn_x = F.relu(self.conv2(cnn_x))\n        cnn_x = self.pool(cnn_x)\n        cnn_x = cnn_x.view(batch_size, -1)\n        \n        # LSTM branch\n        lstm_x = x.unsqueeze(1)\n        lstm_out, _ = self.lstm(lstm_x)\n        lstm_x = lstm_out[:, -1, :]\n        \n        # Concatenate features\n        combined = torch.cat([cnn_x, lstm_x], dim=1)\n        \n        # Classification\n        x = F.relu(self.fc1(combined))\n        x = self.dropout(x)\n        x = F.relu(self.fc2(x))\n        x = self.dropout(x)\n        x = self.fc3(x)\n        \n        return x\n\n# ====================== TRAINING FUNCTIONS ======================\ndef train_epoch(model, dataloader, criterion, optimizer, device):\n    \"\"\"Train for one epoch\"\"\"\n    model.train()\n    losses = []\n    \n    for embeddings, labels in tqdm(dataloader, desc=\"Training\"):\n        embeddings, labels = embeddings.to(device), labels.to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(embeddings)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n        \n        losses.append(loss.item())\n    \n    return np.mean(losses)\n\ndef validate_epoch(model, dataloader, criterion, metric, device):\n    \"\"\"Validate for one epoch\"\"\"\n    model.eval()\n    losses = []\n    scores = []\n    \n    with torch.no_grad():\n        for embeddings, labels in dataloader:\n            embeddings, labels = embeddings.to(device), labels.to(device)\n            \n            outputs = model(embeddings)\n            loss = criterion(outputs, labels)\n            score = metric(torch.sigmoid(outputs), labels.int())\n            \n            losses.append(loss.item())\n            scores.append(score.item())\n    \n    return np.mean(losses), np.mean(scores)\n\ndef train_model(embeddings_source=\"ESM2\", model_type=\"hybrid\", train_ratio=0.9):\n    \"\"\"Complete training pipeline\"\"\"\n    \n    print(f\"\\nTraining {model_type} model with {embeddings_source} embeddings...\")\n    \n    # Create dataset\n    dataset = ProteinDataset(\"train\", embeddings_source, config)\n    \n    # Split into train and validation\n    train_size = int(len(dataset) * train_ratio)\n    val_size = len(dataset) - train_size\n    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n    \n    # Create dataloaders\n    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=2)\n    val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False, num_workers=2)\n    \n    # Initialize model\n    input_dim = config.embeds_dim[embeddings_source]\n    \n    if model_type == \"mlp\":\n        model = ImprovedMLP(input_dim, config.num_labels, config.dropout_rate)\n    elif model_type == \"cnn\":\n        model = AttentionCNN(input_dim, config.num_labels, config.dropout_rate)\n    elif model_type == \"hybrid\":\n        model = HybridModel(input_dim, config.num_labels, config.dropout_rate)\n    else:\n        raise ValueError(f\"Unknown model type: {model_type}\")\n    \n    model = model.to(config.device)\n    \n    # Loss and optimizer\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr, weight_decay=1e-5)\n    scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=config.patience)\n    metric = MultilabelF1Score(num_labels=config.num_labels, average='micro').to(config.device)\n    \n    # Training loop\n    best_val_score = 0\n    best_model_state = None\n    train_losses, val_losses, val_scores = [], [], []\n    \n    for epoch in range(config.n_epochs):\n        print(f\"\\nEpoch {epoch + 1}/{config.n_epochs}\")\n        \n        # Train\n        train_loss = train_epoch(model, train_loader, criterion, optimizer, config.device)\n        train_losses.append(train_loss)\n        \n        # Validate\n        val_loss, val_score = validate_epoch(model, val_loader, criterion, metric, config.device)\n        val_losses.append(val_loss)\n        val_scores.append(val_score)\n        \n        # Learning rate scheduling\n        scheduler.step(val_loss)\n        \n        # Save best model\n        if val_score > best_val_score:\n            best_val_score = val_score\n            best_model_state = model.state_dict().copy()\n        \n        print(f\"Train Loss: {train_loss:.4f}\")\n        print(f\"Val Loss: {val_loss:.4f}, Val F1: {val_score:.4f}\")\n    \n    # Load best model\n    if best_model_state is not None:\n        model.load_state_dict(best_model_state)\n    \n    return model, {\n        'train_losses': train_losses,\n        'val_losses': val_losses,\n        'val_scores': val_scores,\n        'best_score': best_val_score\n    }\n\n# ====================== PREDICTION ======================\ndef predict(model, embeddings_source=\"ESM2\"):\n    \"\"\"Generate predictions for test set\"\"\"\n    \n    print(\"\\nGenerating predictions...\")\n    \n    # Load test dataset\n    test_dataset = ProteinDataset(\"test\", embeddings_source, config)\n    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)\n    \n    # Get label names\n    labels_df = pd.read_csv(config.train_labels_path, sep=\"\\t\", names=[\"EntryID\", \"term\", \"aspect\"])\n    top_terms = labels_df.groupby(\"term\")[\"EntryID\"].count().sort_values(ascending=False)\n    label_names = top_terms[:config.num_labels].index.tolist()\n    \n    # Generate predictions\n    model.eval()\n    predictions = []\n    \n    with torch.no_grad():\n        for embeddings, protein_id in tqdm(test_loader, desc=\"Predicting\"):\n            embeddings = embeddings.to(config.device)\n            outputs = torch.sigmoid(model(embeddings)).cpu().numpy().squeeze()\n            \n            for i, conf in enumerate(outputs):\n                if conf > 0.01:  # Only include predictions above threshold\n                    predictions.append({\n                        'Id': protein_id[0],\n                        'GO term': label_names[i],\n                        'Confidence': min(conf, 0.999)  # Cap at 0.999\n                    })\n    \n    return pd.DataFrame(predictions)\n\n# ====================== ENSEMBLE ======================\ndef ensemble_predictions(predictions_list, weights=None):\n    \"\"\"Combine multiple predictions using weighted averaging\"\"\"\n    \n    if weights is None:\n        weights = [1/len(predictions_list)] * len(predictions_list)\n    \n    # Combine all predictions\n    combined = pd.concat(predictions_list, ignore_index=True)\n    \n    # Group by Id and GO term, taking weighted average\n    ensemble = combined.groupby(['Id', 'GO term'])['Confidence'].apply(\n        lambda x: np.average(x, weights=weights[:len(x)])\n    ).reset_index()\n    \n    return ensemble\n\n# ====================== MAIN EXECUTION ======================\ndef main():\n    \"\"\"Main execution function\"\"\"\n    \n    print(\"Starting CAFA 6 Protein Function Prediction\")\n    print(\"=\" * 60)\n    \n    # Train models with different architectures and embeddings\n    models_and_results = []\n    \n    # Train with ESM2 embeddings\n    if os.path.exists(f\"/kaggle/input/{config.embeds_map['ESM2']}/train_embeddings.npy\"):\n        model_esm2, results_esm2 = train_model(\"ESM2\", \"hybrid\")\n        models_and_results.append((\"ESM2\", model_esm2, results_esm2))\n    \n    # Generate predictions\n    all_predictions = []\n    \n    for embed_source, model, results in models_and_results:\n        print(f\"\\nBest validation F1 for {embed_source}: {results['best_score']:.4f}\")\n        preds = predict(model, embed_source)\n        all_predictions.append(preds)\n    \n    # Combine predictions if multiple models\n    if len(all_predictions) > 1:\n        final_predictions = ensemble_predictions(all_predictions)\n    else:\n        final_predictions = all_predictions[0]\n    \n    # Load and merge with existing predictions if available\n    existing_pred_path = '/kaggle/input/blast-quick-sprof-zero-pred/submission.tsv'\n    if os.path.exists(existing_pred_path):\n        print(\"\\nMerging with existing predictions...\")\n        existing = pd.read_csv(existing_pred_path, sep='\\t', header=None, \n                              names=['Id', 'GO term', 'Confidence'])\n        \n        # Merge predictions\n        merged = pd.merge(existing, final_predictions, on=['Id', 'GO term'], \n                         how='outer', suffixes=('_existing', '_new'))\n        \n        # Combine confidences (take maximum or average)\n        merged['Confidence'] = merged[['Confidence_existing', 'Confidence_new']].max(axis=1)\n        final_predictions = merged[['Id', 'GO term', 'Confidence']]\n    \n    # Save submission\n    print(\"\\nSaving submission file...\")\n    final_predictions.to_csv('submission.tsv', sep='\\t', header=False, index=False)\n    print(f\"Submission saved with {len(final_predictions)} predictions\")\n    \n    # Visualize training results\n    if models_and_results:\n        plt.figure(figsize=(12, 4))\n        \n        plt.subplot(1, 2, 1)\n        for embed_source, _, results in models_and_results:\n            plt.plot(results['val_losses'], label=f\"{embed_source}\")\n        plt.title('Validation Loss')\n        plt.xlabel('Epoch')\n        plt.ylabel('Loss')\n        plt.legend()\n        \n        plt.subplot(1, 2, 2)\n        for embed_source, _, results in models_and_results:\n            plt.plot(results['val_scores'], label=f\"{embed_source}\")\n        plt.title('Validation F1 Score')\n        plt.xlabel('Epoch')\n        plt.ylabel('F1 Score')\n        plt.legend()\n        \n        plt.tight_layout()\n        plt.savefig('training_results.png')\n        plt.show()\n    \n    print(\"\\nDone! Submission file created: submission.tsv\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}