{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14495813,"sourceType":"datasetVersion","datasetId":9258613}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"<!-- ===================== APPROACH HEADER ===================== -->\n<h1 id=\"approach\" style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 24px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n🧠 Approach: ProtLearn-Features\n</h1>\n\n<!-- ===================== MAIN CONTAINER ===================== -->\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p>\n    The central idea of this approach is to represent protein sequences using\n    <strong>ProtLearn</strong>, a bioinformatics-focused Python library designed for extracting\n    features from amino acid sequences.\n  </p>\n\n  <p>\n    ProtLearn implements a collection of well-established protein descriptors that encode\n    sequence composition, local residue patterns, physicochemical properties, and positional\n    distributions. Each descriptor family captures a different aspect of protein structure\n    and function, allowing raw sequences to be transformed into meaningful numerical vectors.\n  </p>\n\n  <p>\n    In this notebook, multiple ProtLearn feature groups are computed independently and\n    concatenated to form a single <strong>695-dimensional representation</strong> per protein.\n    These representations serve as the input to a multi-label classifier that\n    predicts Gene Ontology (GO) terms.\n  </p>\n\n\n  <p style=\"margin-bottom: 0;\">\n    📘 <strong>ProtLearn Documentation:</strong>\n    <a href=\"https://protlearn.readthedocs.io/en/latest/\"\n       style=\"color: #2c5364; font-weight: bold; text-decoration: none;\">\n      https://protlearn.readthedocs.io/\n    </a><br>\n    🐙 <strong>ProtLearn GitHub:</strong>\n    <a href=\"https://github.com/tadorfer/protlearn\"\n       style=\"color: #2c5364; font-weight: bold; text-decoration: none;\">\n      https://github.com/taylor-joren/protlearn\n    </a>\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"!pip install protlearn obonet","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2026-01-14T19:23:24.737047Z","iopub.execute_input":"2026-01-14T19:23:24.737481Z","iopub.status.idle":"2026-01-14T19:23:31.885552Z","shell.execute_reply.started":"2026-01-14T19:23:24.737443Z","shell.execute_reply":"2026-01-14T19:23:31.884599Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<h1 id=\"import\"\n    style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 12px 24px;\n  text-align: center;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n  1. Importing Libraries\n</h1>","metadata":{}},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nfrom Bio import SeqIO\nfrom collections import Counter\nimport gc\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.linear_model import LogisticRegression\n\nfrom protlearn.features import (\n    length, aac, ngram, entropy, \n    ctdc, ctdt, ctdd,\n    paac, apaac,\n    moreau_broto, moran, geary\n)\nfrom protlearn.preprocessing import remove_unnatural, remove_duplicates\n\nimport obonet\nimport torch\nfrom torch import nn\nimport networkx as nx\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom torchmetrics.classification import MultilabelF1Score\n\nfrom matplotlib import pyplot as plt\n\nwarnings.filterwarnings(\"ignore\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:23:31.886887Z","iopub.execute_input":"2026-01-14T19:23:31.887199Z","iopub.status.idle":"2026-01-14T19:23:45.562812Z","shell.execute_reply.started":"2026-01-14T19:23:31.887171Z","shell.execute_reply":"2026-01-14T19:23:45.562106Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<h1 id=\"configuration\"\n    style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 12px 24px;\n  text-align: center;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n  2. Configurations\n</h1>","metadata":{}},{"cell_type":"code","source":"class Config:\n    IA_path  = \"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\"\n    train_terms_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\n    train_taxonomy = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\"\n    train_sequences = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\n    obo_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n    test_sequences = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\n    test_features = \"/kaggle/input/protlearn-features/X_test_data.pkl\"\n    train_features = \"/kaggle/input/protlearn-features/X_train_data.pkl\"\n\n    # Model architecture\n    input_dim = 695    \n    num_labels = 500     \n    hidden_dims = [1024, 512, 512]  \n    dropout = 0.3         \n\n    # Training parameters\n    batch_size = 128\n    train_split = 0.9    \n    lr = 1e-3             \n    epochs = 100\n    patience = 15         \n\n    # Device configuration\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    seed = 42\n\nconfig = Config()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:23:45.569421Z","iopub.execute_input":"2026-01-14T19:23:45.569725Z","iopub.status.idle":"2026-01-14T19:23:45.650549Z","shell.execute_reply.started":"2026-01-14T19:23:45.569694Z","shell.execute_reply":"2026-01-14T19:23:45.649665Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== SEQUENCE PREPROCESSING ===================== -->\n<h1 id=\"sequence-preprocessing\" style=\"\n  background: linear-gradient(135deg, #1c2b36, #2c3e50, #34495e);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n3. Sequence Loading and Cleaning\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p>\n    Protein sequences are loaded from FASTA files and converted into a pandas DataFrame\n    containing UniProt entry IDs and raw amino acid sequences.\n  </p>\n\n  <p>\n    ProtLearn only accepts sequences composed of the <strong>20 standard amino acids</strong>.\n    However, the dataset includes some non-standard or ambiguous residues that must be\n    resolved prior to feature extraction.\n  </p>\n\n  <p>\n    To ensure valid ProtLearn input, the following cleaning rules are applied:\n  </p>\n\n  <ul>\n    <li>\n      <strong>U → C</strong>\n      (selenocysteine — the rare <em>21st amino acid</em> — mapped to cysteine due to chemical similarity)\n    </li>\n    <li>\n      <strong>B → D</strong>\n      (ambiguous residue representing either aspartic acid or asparagine, resolved to aspartic acid)\n    </li>\n    <li>\n      <strong>Z → E</strong>\n      (ambiguous residue representing either glutamic acid or glutamine, resolved to glutamic acid)\n    </li>\n    <li>\n      <strong>X → removed</strong>\n      (unknown or masked amino acid with no biological specificity)\n    </li>\n  </ul>\n\n  <p style=\"margin-bottom: 0;\">\n    After cleaning, all sequences strictly follow the standard amino acid alphabet and are\n    fully compatible with ProtLearn feature extraction.\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"def load_fasta_as_dataframe(fasta_path: str, train_data:bool = True):\n    ids,sequences = [],[]\n\n    for record in SeqIO.parse(fasta_path, \"fasta\"):\n        if train_data:\n            uniprot_id = record.id.split(\"|\")[1]\n        else:\n            uniprot_id = record.id\n            \n        ids.append(uniprot_id)\n        sequences.append(str(record.seq))\n\n    return pd.DataFrame({\n        \"EntryID\": ids,\n        \"sequences\": sequences\n    })\n    \n\n#Handling non standard Amino acid residues\n\ndef clean_sequence(seq):\n    return (\n        seq.replace(\"U\", \"C\")   # selenocysteine → cysteine\n           .replace(\"B\", \"D\")   # Asp/Asn → Asp\n           .replace(\"Z\", \"E\")   # Glu/Gln → Glu\n           .replace(\"X\", \"\")    # remove unknowns\n    )\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:23:45.651506Z","iopub.execute_input":"2026-01-14T19:23:45.651790Z","iopub.status.idle":"2026-01-14T19:23:45.670109Z","shell.execute_reply.started":"2026-01-14T19:23:45.651768Z","shell.execute_reply":"2026-01-14T19:23:45.669356Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== FEATURE EXTRACTION ===================== -->\n<h1 id=\"feature-extraction\" style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n4. Feature Extraction with ProtLearn\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p>\n    Cleaned protein sequences are converted into fixed-length numerical representations\n    using <strong>ProtLearn</strong>, a feature-engineering library for protein sequence analysis.\n  </p>\n\n  <p>\n    For each protein, multiple complementary feature groups are extracted and concatenated\n    into a single vector.\n  </p>\n\n  <ul>\n    <li>\n      <strong>Sequence Length (1):</strong>\n      Total number of amino acids in the protein.\n    </li>\n    <li>\n      <strong>Amino Acid Composition – AAC (20):</strong>\n      Relative frequency of each standard amino acid.\n    </li>\n    <li>\n      <strong>Dipeptide Composition (400):</strong>\n      Frequency of all possible adjacent amino acid pairs.\n    </li>\n    <li>\n      <strong>Shannon Entropy (1):</strong>\n      Measures sequence complexity and variability.\n    </li>\n    <li>\n      <strong>CTD Features:</strong>\n      Physicochemical descriptors capturing amino acid\n      <em>composition</em>, <em>transition</em>, and <em>distribution</em>\n      along the sequence.\n    </li>\n  </ul>\n\n  <p>\n    All feature groups are horizontally concatenated, resulting in a\n    <strong>695-dimensional feature vector per protein</strong>.\n  </p>\n\n\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"def extract_features(sequences):\n    \"\"\"\n    Extract  feature set from protein sequences.\n    \n    Args:\n        sequences: List of protein sequences (strings)\n    \n    Returns:\n        numpy array of shape (n_proteins, 695) containing all features\n    \"\"\"\n    \n    # 1. Length - basic sequence length (1 feature)\n    X_len = length(sequences)\n    \n    # 2. Amino Acid Composition - frequency of each amino acid (20 features)\n    X_aac, _ = aac(sequences, method='relative')\n    \n    # 3. Dipeptide Composition - frequency of amino acid pairs (400 features)\n    X_di, _ = ngram(sequences, n=2)\n    \n    # 4. Shannon Entropy - measure of sequence complexity (1 feature)\n    X_entropy = entropy(sequences)\n    \n    # 5. Composition-Transition-Distribution (CTD) Features\n    # These capture physicochemical properties of amino acids\n    X_ctdc, _ = ctdc(sequences)  # Composition\n    X_ctdt, _ = ctdt(sequences)  # Transition between property groups\n    X_ctdd, _ = ctdd(sequences)  # Distribution of properties along sequence\n    \n    # Combine all features horizontally\n    features = [X_len, X_aac, X_di, X_entropy, X_ctdc, X_ctdt, X_ctdd]\n    X_all = np.hstack(features)\n    \n    print(f\"\\n✓ Feature extraction complete. Shape: {X_all.shape}\")\n    print(f\"  Total features per protein: {X_all.shape[1]}\")\n    \n    return X_all\n\n\n# # Load sequences\n# train_sequences_df = load_fasta_as_dataframe(config.train_sequences, train_data=True)\n# test_sequences_df = load_fasta_as_dataframe(config.test_sequences, train_data=False)\n\n# # Clean sequences\n# train_sequences_df[\"sequences\"] = train_sequences_df[\"sequences\"].apply(clean_sequence)\n# test_sequences_df[\"sequences\"] = test_sequences_df[\"sequences\"].apply(clean_sequence)\n\n# # Extract features\n# train_sequences_list = train_sequences_df['sequences'].tolist()\n# test_sequences_list = test_sequences_df['sequences'].tolist()\n\n# X_train = extract_features(train_sequences_list)\n# X_test = extract_features(test_sequences_list)\n\n\n# print(\"✓ Features saved successfully!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:23:45.671006Z","iopub.execute_input":"2026-01-14T19:23:45.671334Z","iopub.status.idle":"2026-01-14T19:23:45.685991Z","shell.execute_reply.started":"2026-01-14T19:23:45.671307Z","shell.execute_reply":"2026-01-14T19:23:45.685247Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== FEATURE & LABEL PREPARATION ===================== -->\n<h1 id=\"feature-label-preparation\" style=\"\n  background: linear-gradient(135deg, #16222a, #3a6073);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n5. Feature and Label Preparation\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p style=\"margin-bottom: 0;\">\n    Precomputed ProtLearn features are loaded from disk and aligned using UniProt\n    EntryIDs, while Gene Ontology (GO) annotations are grouped per protein to form\n    multi-label targets. To account for the hierarchical nature of GO, each annotated\n    term is expanded to include all its ancestor terms using the GO graph, after which\n    only the top N most frequent GO terms are retained. These propagated label sets are\n    then binarized into fixed-length target vectors and merged with the ProtLearn\n    feature matrix to construct the final training dataset,\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"train_features_df = pd.read_pickle(config.train_features).rename({\"ID\":\"EntryID\"}, axis=1)\ntest_features_df = pd.read_pickle(config.test_features).rename({\"ID\":\"EntryID\"}, axis=1)\n\nprint(f\"train_features_df shape: {train_features_df.shape}\")\nprint(f\"test_features_df shape:  {test_features_df.shape}\")\n\n# Load GO term labels\ntrain_terms_df = pd.read_csv(config.train_terms_path, sep=\"\\t\")\n\n# Group GO terms by protein ID\nlabels_df = (\n    train_terms_df\n    .groupby(\"EntryID\")[\"term\"]\n    .apply(list)\n    .reset_index(name=\"labels\")\n)\n\n\ngo_graph = obonet.read_obo(config.obo_path)\n\ndef get_ancestors(graph, term):\n    if term not in graph:\n        return set()\n    return nx.descendants(graph, term) | {term}\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"PROPAGATING GO TERM ANCESTORS\")\nprint(\"=\"*60)\n\n# Select top N most frequent GO terms\ntop_go = (\n    train_terms_df[\"term\"]\n    .value_counts()\n    .iloc[:config.num_labels]\n    .index\n    .tolist()\n)\ntop_go_set = set(top_go)\n\n# Group GO terms by protein ID\nprotein_to_terms = train_terms_df.groupby(\"EntryID\")[\"term\"].apply(set).to_dict()\n\n# Propagate ancestors for each protein\npropagated_labels = {}\nprint(\"Propagating ancestors for each protein...\")\n\nfor protein_id, terms in tqdm(protein_to_terms.items(), desc=\"Propagating\"):\n    all_ancestors = set()\n    for term in terms:\n        # Get all ancestors for this term\n        ancestors = get_ancestors(go_graph, term)\n        all_ancestors.update(ancestors)\n    \n    # Filter to only keep terms in our top N list\n    propagated_labels[protein_id] = all_ancestors.intersection(top_go_set)\n\n# Create labels DataFrame with propagated terms\nlabels_data = []\nfor protein_id, terms in propagated_labels.items():\n    if terms:  # Only include proteins with at least one term\n        labels_data.append({\n            \"EntryID\": protein_id,\n            \"labels\": list(terms)\n        })\n\nlabels_df = pd.DataFrame(labels_data)\n\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"PREPROCESSING DATA\")\nprint(\"=\"*60)\n\n# Initialize MultiLabelBinarizer for one-hot encoding GO terms\nmlb = MultiLabelBinarizer(classes=top_go)\nmlb.fit([top_go])\n\n# Convert GO term lists to binary target vectors\nlabels_df[\"targets\"] = list(mlb.transform(labels_df[\"labels\"]))\n\n# Create training DataFrame with embeddings and labels\n\n\ntrain_df = train_features_df.merge(labels_df, on=\"EntryID\", how=\"inner\")\n\n# Create test DataFrame\ntest_df = test_features_df \n\nprint(f\"Training data shape: {train_df.shape}\")\nprint(f\"Test data shape: {test_df.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:23:54.671607Z","iopub.execute_input":"2026-01-14T19:23:54.671948Z","iopub.status.idle":"2026-01-14T19:24:13.935633Z","shell.execute_reply.started":"2026-01-14T19:23:54.671895Z","shell.execute_reply":"2026-01-14T19:24:13.934775Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== MODEL ARCHITECTURE ===================== -->\n<h1 id=\"model-architecture\" style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n6. Model Architecture\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p style=\"margin-bottom: 0;\">\n    The model is a fully connected multi-layer perceptron (MLP) designed for\n    multi-label protein function prediction. It consists of stacked hidden layers\n    with batch normalization, ReLU activations, and dropout for regularization.\n    Residual connections are added every two layers to facilitate gradient flow\n    and improve training stability. The output layer maps to the fixed set of\n    selected Gene Ontology terms, producing one score per label for each protein.\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"class MLP(nn.Module):\n    \n    def __init__(self, input_dim=1024, num_labels=500, hidden_dims=[1024, 512, 512], dropout=0.3):\n        super().__init__()\n        self.layers = nn.ModuleList()\n        self.bns = nn.ModuleList()\n        self.dropouts = nn.ModuleList()\n\n        # Build hidden layers\n        prev_dim = input_dim\n        for h in hidden_dims:\n            self.layers.append(nn.Linear(prev_dim, h))\n            self.bns.append(nn.BatchNorm1d(h))\n            self.dropouts.append(nn.Dropout(dropout))\n            prev_dim = h\n\n        # Output layer\n        self.out = nn.Linear(prev_dim, num_labels)\n\n    def forward(self, x):\n        residual = x\n        \n        # Pass through hidden layers with residual connections\n        for i, layer in enumerate(self.layers):\n            x = layer(x)\n            x = self.bns[i](x)\n            x = torch.relu(x)\n            x = self.dropouts[i](x)\n            \n            # Add residual connection every 2 layers\n            if i % 2 == 1:\n                if residual.shape[1] == x.shape[1]:\n                    x = x + residual\n                residual = x\n            else:\n                residual = x\n        \n        # Output layer\n        x = self.out(x)\n        return x","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:24:13.936858Z","iopub.execute_input":"2026-01-14T19:24:13.937413Z","iopub.status.idle":"2026-01-14T19:24:13.944614Z","shell.execute_reply.started":"2026-01-14T19:24:13.937388Z","shell.execute_reply":"2026-01-14T19:24:13.943962Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== DATASET & DATALOADERS ===================== -->\n<h1 id=\"dataset-dataloaders\" style=\"\n  background: linear-gradient(135deg, #16222a, #3a6073);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n7. Dataset and Data Loading\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p style=\"margin-bottom: 0;\">\n    A custom PyTorch dataset is defined to serve ProtLearn feature vectors and\n    multi-label GO targets, returning feature–label pairs during training and\n    feature–ID pairs during inference. The full training set is split into\n    training and validation subsets using a fixed random seed to ensure\n    reproducibility, and DataLoaders are constructed for efficient mini-batch\n    processing with shuffling applied only to the training data.\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"class ProteinFunctionDataset(Dataset):\n    \"\"\"Custom Dataset for protein embeddings and GO term labels\"\"\"\n    \n    def __init__(self, df, train=True):\n        self.df = df.reset_index(drop=True)\n        self.train = train\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        embed = torch.tensor(self.df.loc[idx, \"features\"], dtype=torch.float32)\n        \n        if self.train:\n            target = torch.tensor(self.df.loc[idx, \"targets\"], dtype=torch.float32)\n            return embed, target\n        else:\n            entry_id = self.df.loc[idx, \"EntryID\"]\n            return embed, entry_id\n\n# Create datasets\nfull_train_dataset = ProteinFunctionDataset(train_df, train=True)\ntest_dataset = ProteinFunctionDataset(test_df, train=False)\n\n# Split training data into train and validation sets\ntrain_size = int(len(full_train_dataset) * config.train_split)\nval_size = len(full_train_dataset) - train_size\n\ng = torch.Generator().manual_seed(config.seed)\ntrain_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size], generator=g)\n\n# Create dataloaders\ntrain_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)\nval_loader   = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False)\ntest_loader  = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)\n\nprint(f\"\\nTrain batches: {len(train_loader)}\")\nprint(f\"Validation batches: {len(val_loader)}\")\nprint(f\"Test batches: {len(test_loader)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:24:13.945605Z","iopub.execute_input":"2026-01-14T19:24:13.945863Z","iopub.status.idle":"2026-01-14T19:24:14.020145Z","shell.execute_reply.started":"2026-01-14T19:24:13.945844Z","shell.execute_reply":"2026-01-14T19:24:14.019448Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== MODEL TRAINING ===================== -->\n<h1 id=\"model-training\" style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n8. Model Training\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p style=\"margin-bottom: 0;\">\n    The MLP model is trained using mini-batch gradient descent with the Adam optimizer,\n    applying cross-entropy loss for multi-label classification. Training proceeds over\n    multiple epochs, with performance evaluated on a validation set at each epoch using\n    F1 score. A learning rate scheduler adapts the optimizer based on validation loss, and\n    early stopping is implemented to save the model checkpoint corresponding to the best\n    validation performance, preventing overfitting and ensuring optimal generalization.\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"def train_model(model, train_loader, val_loader, cfg, save_path=\"best_model.pt\"):\n  \n    \n    model.to(config.device)\n\n    # Initialize optimizer, scheduler, and loss function\n    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer, mode=\"min\", factor=0.5, patience=2\n    )\n    loss_fn = nn.CrossEntropyLoss()\n    # loss_fn = nn.BCEWithLogitsLoss() \n    f1_score = MultilabelF1Score(num_labels=config.num_labels).to(config.device)\n\n    # Tracking metrics\n    history = {\n        \"train_loss\": [],\n        \"val_loss\": [],\n        \"train_f1\": [],\n        \"val_f1\": []\n    }\n\n    best_val_loss = float(\"inf\")\n    patience_counter = 0\n\n    print(\"\\n\" + \"=\"*60)\n    print(\"TRAINING\")\n    print(\"=\"*60)\n\n    for epoch in range(config.epochs):\n        # ===== TRAINING PHASE =====\n        model.train()\n        train_losses, train_f1s = [], []\n\n        for embeds, targets in train_loader:\n            embeds = embeds.to(config.device)\n            targets = targets.to(config.device)\n\n            # Forward pass\n            optimizer.zero_grad()\n            logits = model(embeds)\n            loss = loss_fn(logits, targets)\n            \n            # Backward pass\n            loss.backward()\n            optimizer.step()\n\n            # Calculate metrics\n            f1 = f1_score(logits, targets)\n            train_losses.append(loss.item())\n            train_f1s.append(f1.item())\n\n        avg_train_loss = np.mean(train_losses)\n        avg_train_f1 = np.mean(train_f1s)\n\n        # ===== VALIDATION PHASE =====\n        model.eval()\n        val_losses, val_f1s = [], []\n\n        with torch.no_grad():\n            for embeds, targets in val_loader:\n                embeds = embeds.to(config.device)\n                targets = targets.to(config.device)\n\n                logits = model(embeds)\n                loss = loss_fn(logits, targets)\n\n                f1 = f1_score(logits, targets)\n                val_losses.append(loss.item())\n                val_f1s.append(f1.item())\n\n        avg_val_loss = np.mean(val_losses)\n        avg_val_f1 = np.mean(val_f1s)\n\n        # Save metrics\n        history[\"train_loss\"].append(avg_train_loss)\n        history[\"val_loss\"].append(avg_val_loss)\n        history[\"train_f1\"].append(avg_train_f1)\n        history[\"val_f1\"].append(avg_val_f1)\n\n        # Print progress\n        print(\n            f\"Epoch {epoch+1:3d}/{cfg.epochs} | \"\n            f\"Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f} | \"\n            f\"Train F1: {avg_train_f1:.4f} | Val F1: {avg_val_f1:.4f}\"\n        )\n\n        # Learning rate scheduling\n        scheduler.step(avg_val_loss)\n\n        # ===== CHECKPOINTING & EARLY STOPPING =====\n        if avg_val_loss < best_val_loss:\n            best_val_loss = avg_val_loss\n            patience_counter = 0\n            torch.save(model.state_dict(), save_path)\n            print(f\"💾 Best model saved at epoch {epoch+1}\")\n        else:\n            patience_counter += 1\n            print(f\"⏳ No improvement. Patience: {patience_counter}/{config.patience}\")\n\n        if patience_counter >= cfg.patience:\n            print(\"🛑 Early stopping triggered.\")\n            break\n\n    print(f\"\\n✅ Training complete! Best Val Loss: {best_val_loss:.4f}\")\n    return history\n\n\nmodel = MLP(\n    input_dim=config.input_dim,\n    num_labels=config.num_labels,\n    hidden_dims=config.hidden_dims,\n    dropout=config.dropout\n)\n\nhistory = train_model(model, train_loader, val_loader, config, save_path=\"best_model.pt\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:24:14.021028Z","iopub.execute_input":"2026-01-14T19:24:14.021319Z","iopub.status.idle":"2026-01-14T19:32:37.599250Z","shell.execute_reply.started":"2026-01-14T19:24:14.021297Z","shell.execute_reply":"2026-01-14T19:32:37.598315Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== TRAINING VISUALIZATION ===================== -->\n<h1 id=\"training-visualization\" style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n9. Training Visualization\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p style=\"margin-bottom: 0;\">\n    Training progress is visualized by plotting both loss and F1 score over epochs\n    for the training and validation sets. These plots provide insight into model\n    convergence, performance trends, and potential overfitting. \n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"# VISUALIZATION\ndef plot_training_history(history, save_path=\"training_history.png\"):\n    \"\"\"Plot training and validation metrics\"\"\"\n    \n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n    \n    epochs = range(1, len(history[\"train_loss\"]) + 1)\n    \n    # Loss plot\n    ax1.plot(epochs, history[\"train_loss\"], 'b-', label=\"Train Loss\", linewidth=2)\n    ax1.plot(epochs, history[\"val_loss\"], 'r-', label=\"Val Loss\", linewidth=2)\n    ax1.set_xlabel(\"Epoch\", fontsize=12)\n    ax1.set_ylabel(\"Loss\", fontsize=12)\n    ax1.set_title(\"Training and Validation Loss\", fontsize=14, fontweight='bold')\n    ax1.legend(fontsize=10)\n    ax1.grid(True, alpha=0.3)\n    \n    # F1 score plot\n    ax2.plot(epochs, history[\"train_f1\"], 'b-', label=\"Train F1\", linewidth=2)\n    ax2.plot(epochs, history[\"val_f1\"], 'r-', label=\"Val F1\", linewidth=2)\n    ax2.set_xlabel(\"Epoch\", fontsize=12)\n    ax2.set_ylabel(\"F1 Score\", fontsize=12)\n    ax2.set_title(\"Training and Validation F1 Score\", fontsize=14, fontweight='bold')\n    ax2.legend(fontsize=10)\n    ax2.grid(True, alpha=0.3)\n    \n    plt.tight_layout()\n    plt.savefig(save_path, dpi=300, bbox_inches='tight')\n    print(f\"\\n📊 Training history plot saved to {save_path}\")\n    plt.show()\n\n\n\n# Plot training history\nplot_training_history(history)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:32:37.600394Z","iopub.execute_input":"2026-01-14T19:32:37.600638Z","iopub.status.idle":"2026-01-14T19:32:38.856674Z","shell.execute_reply.started":"2026-01-14T19:32:37.600618Z","shell.execute_reply":"2026-01-14T19:32:38.855994Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<!-- ===================== INFERENCE & SUBMISSION ===================== -->\n<h1 id=\"inference-submission\" style=\"\n  background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);\n  color: white;\n  padding: 14px 26px;\n  border-radius: 12px;\n  font-size: 22px;\n  font-weight: bold;\n  text-align: center;\n  font-family: Arial, sans-serif;\n  margin: 20px 0;\n\">\n10. Inference and Submission Generation\n</h1>\n\n<div style=\"\n  background: white;\n  border-radius: 12px;\n  border: 1px solid #e0e0e0;\n  font-family: Arial, sans-serif;\n  color: #333;\n  padding: 20px;\n  font-size: 15px;\n  line-height: 1.6;\n  margin-bottom: 25px;\n\">\n\n  <p style=\"margin-bottom: 0;\">\n    The trained MLP model is loaded from the best checkpoint and used to predict\n    Gene Ontology term probabilities for all test proteins. Predictions are\n    computed in mini-batches, with a sigmoid applied to produce probabilities\n    for each GO label. For each protein, only terms with positive probability\n    are retained, sorted by confidence, and truncated to a maximum per-protein\n    limit. Probabilities are formatted to three significant figures, and the\n    resulting predictions are saved in tab-separated format for submission.\n  </p>\n\n</div>\n","metadata":{}},{"cell_type":"code","source":"# Load the best model checkpoint saved during training\nmodel.load_state_dict(torch.load(\"best_model.pt\", map_location=config.device))\nprint(\"✅ Best model loaded successfully!\")\n\ndef generate_submission(\n    model,\n    test_loader,\n    go_terms,\n    device,\n    output_path=\"submission.tsv\",\n    max_terms_per_protein=1500,\n    min_prob=1e-6\n):\n   \n    \n    model.eval()\n    \n    all_ids = []\n    all_probs = []\n\n    # Generate predictions\n    with torch.no_grad():\n        for embeds, ids in tqdm(test_loader, desc=\"Predicting\"):\n            embeds = embeds.to(device)\n            logits = model(embeds)\n            probs = torch.sigmoid(logits).cpu().numpy()\n\n            all_probs.append(probs)\n            all_ids.extend(ids)\n\n    all_probs = np.vstack(all_probs)  # Shape: (N_test, num_labels)\n    all_ids = np.array(all_ids)\n\n    # Format predictions for submission\n    rows = []\n\n    for i, prot_id in enumerate(all_ids):\n        probs = all_probs[i]\n\n        # Keep only positive probabilities\n        idx = np.where(probs > 0)[0]\n        if len(idx) == 0:\n            continue\n\n        # Sort by confidence (descending)\n        idx = idx[np.argsort(probs[idx])[::-1]]\n\n        # Limit to top 1500 terms per protein \n        idx = idx[:max_terms_per_protein]\n\n        for j in idx:\n            p = probs[j]\n            if p <= 0:\n                continue\n\n            # Format to 3 significant figures \n            p_fmt = float(f\"{p:.3g}\")\n            if p_fmt <= 0:\n                p_fmt = min_prob\n\n            rows.append([prot_id, go_terms[j], p_fmt])\n\n    # Create submission DataFrame\n    sub_df = pd.DataFrame(rows)\n    sub_df.to_csv(output_path, sep=\"\\t\", header=False, index=False)\n\n    print(f\"\\n✅ Submission saved to {output_path}\")\n    print(f\"Total predictions: {len(sub_df):,}\")\n\n    return sub_df\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"GENERATING SUBMISSION\")\nprint(\"=\"*60)\n\nsubmission_df = generate_submission(\n    model=model,\n    test_loader=test_loader,\n    go_terms=top_go,\n    device=config.device,\n    output_path=\"submission.tsv\"\n)\n\n# Display sample predictions\nprint(\"\\nSample predictions:\")\nprint(submission_df.head(10))\n\nprint(\"\\n\" + \"=\"*60)\nprint(\"DONE!\")\nprint(\"=\"*60)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T19:32:38.857756Z","iopub.execute_input":"2026-01-14T19:32:38.858038Z"}},"outputs":[],"execution_count":null}]}