{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-14T15:04:08.703672Z","iopub.execute_input":"2025-12-14T15:04:08.704042Z","iopub.status.idle":"2025-12-14T15:04:08.710528Z","shell.execute_reply.started":"2025-12-14T15:04:08.704015Z","shell.execute_reply":"2025-12-14T15:04:08.709486Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Given protein amino acid sequences and must build a model that predicts what each protein does. The “function” is expressed using GO terms (Gene Ontology), which describe:\n* Molecular Function (MF) – what the protein does at the biochemical level\n* Biological Process (BP) – which larger process it participates in\n* Cellular Component (CC) – where in the cell it operates\n\n\n> Predict the correct GO terms for each protein in the test dataset, with probability scores.\n\n\n*\"Many proteins still have unknown functions, even though we know their sequences. Predicting function automatically will help scientists study diseases and develop new therapies faster.\"*\n\n","metadata":{}},{"cell_type":"markdown","source":"What you need to do (step-by-step)\n1. Understand the prediction target\nYou must predict Gene Ontology (GO) terms across three sub-ontologies:\nMolecular Function (MF) – what the protein does at a biochemical level.\nBiological Process (BP) – what broader biological processes it participates in.\nCellular Component (CC) – where in the cell it operates.\nEach protein can have many GO terms in each category.\nThis is a multi-label classification problem with a large ontology.","metadata":{}},{"cell_type":"markdown","source":"# Setup","metadata":{}},{"cell_type":"code","source":"!pip install obonet Bio\n!pip install biopython pandas scipy numpy","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Explore dataset","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nfrom Bio import SeqIO","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"go-basic.obo\n* Gene Ontology (GO) hierarchy.\nYou need this for:\n* Ancestor propagation\n* Parent–child relationships\n* Filtering invalid GO terms.","metadata":{}},{"cell_type":"code","source":"print(\"=== TRAIN: go-basic.obo (Ontology structure) ===\")\n# Only printing first few lines\nwith open(\"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\", \"r\") as f:\n    for i in range(20):\n        print(f.readline().strip())\nprint(\"... (ontology continues)\")\nprint()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:51:03.942836Z","iopub.execute_input":"2025-11-23T13:51:03.943218Z","iopub.status.idle":"2025-11-23T13:51:03.952570Z","shell.execute_reply.started":"2025-11-23T13:51:03.943197Z","shell.execute_reply":"2025-11-23T13:51:03.951638Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"train_sequences.fasta\n\n* Protein sequences with IDs.\n* Input for your feature extraction / embeddings.","metadata":{}},{"cell_type":"code","source":"print(\"=== TRAIN: train_sequences.fasta ===\")\ntrain_sequences = list(SeqIO.parse(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\", \"fasta\"))\nprint(\"Total sequences:\", len(train_sequences))\nprint(\"Example record:\")\nprint(\"ID:\", train_sequences[0].id)\nprint(\"Sequence length:\", len(train_sequences[0].seq))\nprint()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:39:20.993872Z","iopub.execute_input":"2025-11-23T13:39:20.994210Z","iopub.status.idle":"2025-11-23T13:39:22.999049Z","shell.execute_reply.started":"2025-11-23T13:39:20.994185Z","shell.execute_reply":"2025-11-23T13:39:22.998152Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"train_taxonomy.tsv\n\n* Protein ID → Taxon ID.\n* Useful because function is taxon-dependent.\n* Some models incorporate taxon embeddings.","metadata":{}},{"cell_type":"code","source":"print(\"=== TRAIN: train_taxonomy.tsv ===\")\ntrain_tax = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"taxon_id\"])\nprint(train_tax.head())\nprint(\"Taxon count:\", train_tax.taxon_id.nunique())\nprint()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:45:34.980048Z","iopub.execute_input":"2025-11-23T13:45:34.980327Z","iopub.status.idle":"2025-11-23T13:45:35.032072Z","shell.execute_reply.started":"2025-11-23T13:45:34.980307Z","shell.execute_reply":"2025-11-23T13:45:35.031205Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"train_terms.tsv\n* For each protein: the GO terms it is annotated with.\n* This is your training label file.","metadata":{}},{"cell_type":"code","source":"print(\"=== TRAIN: train_terms.tsv ===\")\ntrain_terms = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"go_term\"])\nprint(train_terms.head())\nprint(\"Total annotations:\", len(train_terms))\nprint(\"Unique GO terms:\", train_terms.go_term.nunique())\nprint()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:44:39.966862Z","iopub.execute_input":"2025-11-23T13:44:39.967504Z","iopub.status.idle":"2025-11-23T13:44:40.396669Z","shell.execute_reply.started":"2025-11-23T13:44:39.967476Z","shell.execute_reply":"2025-11-23T13:44:40.395470Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**IA.tsv**\n* Information Accretion (term information content).\n* Useful to weight rare GO terms during training.","metadata":{}},{"cell_type":"code","source":"print(\"=== TRAIN: IA.tsv (Information Accretion values) ===\")\nia = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\", sep=\"\\t\", header=None, names=[\"go_term\", \"IA_value\"])\nprint(ia.head())\nprint(\"Total GO terms with IA values:\", len(ia))\nprint()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:47:47.472405Z","iopub.execute_input":"2025-11-23T13:47:47.472799Z","iopub.status.idle":"2025-11-23T13:47:47.522641Z","shell.execute_reply.started":"2025-11-23T13:47:47.472778Z","shell.execute_reply":"2025-11-23T13:47:47.521784Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"sample_submission.tsv\n* The correct output format that must generate.","metadata":{}},{"cell_type":"code","source":"print(\"=== TRAIN: sample_submission.tsv ===\")\nsample_sub = pd.read_csv(\n    \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\",\n    sep=\"\\t\",\n    engine=\"python\",\n    comment=\"#\",\n    on_bad_lines=\"skip\"\n)\nprint(sample_sub.head())\nprint(sample_sub.columns)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:54:52.661655Z","iopub.execute_input":"2025-11-23T13:54:52.662045Z","iopub.status.idle":"2025-11-23T13:54:52.745490Z","shell.execute_reply.started":"2025-11-23T13:54:52.662019Z","shell.execute_reply":"2025-11-23T13:54:52.744483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -------------------------------------------------------------------\n# 2. TEST DATA\n# -------------------------------------------------------------------\n\nprint(\"=== TEST: testsuperset.fasta ===\")\ntest_sequences = list(SeqIO.parse(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\", \"fasta\"))\nprint(\"Total test sequences:\", len(test_sequences))\nprint(\"Example record:\")\nprint(\"ID:\", test_sequences[0].id)\nprint(\"Sequence length:\", len(test_sequences[0].seq))\nprint()\n\nprint(\"=== TEST: testsuperset-taxon-list.tsv ===\")\ntest_tax = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"taxon_id\"])\nprint(test_tax.head())\nprint(\"Test taxon count:\", test_tax.taxon_id.nunique())\nprint()\n\nprint(\"=== DONE ===\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:55:57.919076Z","iopub.execute_input":"2025-11-23T13:55:57.919744Z","iopub.status.idle":"2025-11-23T13:56:01.497696Z","shell.execute_reply.started":"2025-11-23T13:55:57.919718Z","shell.execute_reply":"2025-11-23T13:56:01.496868Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nfrom Bio import SeqIO\n\n# -------------------------------------------------------------------\n# 1. TRAINING DATA\n# -------------------------------------------------------------------\n\nprint(\"=== TRAIN: train_sequences.fasta ===\")\ntrain_sequences = list(SeqIO.parse(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\", \"fasta\"))\nprint(\"Total sequences:\", len(train_sequences))\nprint(\"Example record:\")\nprint(\"ID:\", train_sequences[0].id)\nprint(\"Sequence length:\", len(train_sequences[0].seq))\nprint()\n\nprint(\"=== TRAIN: train_terms.tsv ===\")\ntrain_terms = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"go_term\"])\nprint(train_terms.head())\nprint(\"Total annotations:\", len(train_terms))\nprint(\"Unique GO terms:\", train_terms.go_term.nunique())\nprint()\n\nprint(\"=== TRAIN: train_taxonomy.tsv ===\")\ntrain_tax = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"taxon_id\"])\nprint(train_tax.head())\nprint(\"Taxon count:\", train_tax.taxon_id.nunique())\nprint()\n\nprint(\"=== TRAIN: IA.tsv (Information Accretion values) ===\")\nia = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\", sep=\"\\t\", header=None, names=[\"go_term\", \"IA_value\"])\nprint(ia.head())\nprint(\"Total GO terms with IA values:\", len(ia))\nprint()\n\nprint(\"=== TRAIN: go-basic.obo (Ontology structure) ===\")\n# Only printing first few lines\nwith open(\"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\", \"r\") as f:\n    for i in range(20):\n        print(f.readline().strip())\nprint(\"... (ontology continues)\")\nprint()\n\nprint(\"=== TRAIN: sample_submission.tsv ===\")\nsample_sub = pd.read_csv(\n    \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\",\n    sep=\"\\t\",\n    engine=\"python\",\n    comment=\"#\",\n    on_bad_lines=\"skip\"\n)\nprint(sample_sub.head())\nprint(sample_sub.columns)\n\n\n# -------------------------------------------------------------------\n# 2. TEST DATA\n# -------------------------------------------------------------------\n\nprint(\"=== TEST: testsuperset.fasta ===\")\ntest_sequences = list(SeqIO.parse(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\", \"fasta\"))\nprint(\"Total test sequences:\", len(test_sequences))\nprint(\"Example record:\")\nprint(\"ID:\", test_sequences[0].id)\nprint(\"Sequence length:\", len(test_sequences[0].seq))\nprint()\n\nprint(\"=== TEST: testsuperset-taxon-list.tsv ===\")\ntest_tax = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\", sep=\"\\t\", header=None, names=[\"protein_id\", \"taxon_id\"])\nprint(test_tax.head())\nprint(\"Test taxon count:\", test_tax.taxon_id.nunique())\nprint()\n\nprint(\"=== DONE ===\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-23T13:56:32.395516Z","iopub.execute_input":"2025-11-23T13:56:32.395837Z","iopub.status.idle":"2025-11-23T13:56:35.662745Z","shell.execute_reply.started":"2025-11-23T13:56:32.395816Z","shell.execute_reply":"2025-11-23T13:56:35.661887Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Code","metadata":{}},{"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader, random_split\nfrom transformers import AutoTokenizer, EsmModel\nfrom Bio import SeqIO\nimport obonet\nimport networkx as nx\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nfrom collections import Counter\nimport gc\n\n# ==========================================\n# CONFIGURATION\n# ==========================================\nclass Config:\n    # Paths (Adjust these if running locally)\n    BASE_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n    TRAIN_FASTA = os.path.join(BASE_DIR, 'Train', 'train_sequences.fasta')\n    TRAIN_TERMS = os.path.join(BASE_DIR, 'Train', 'train_terms.tsv')\n    TEST_FASTA  = os.path.join(BASE_DIR, 'Test', 'testsuperset.fasta')\n    OBO_FILE    = os.path.join(BASE_DIR, 'Train', 'go-basic.obo')\n    \n    # Model Parameters\n    MODEL_NAME = \"facebook/esm2_t6_8M_UR50D\" # Small, fast model for testing. \n                                              # For winning: use \"facebook/esm2_t33_650M_UR50D\"\n    MAX_SEQ_LEN = 1022  # ESM limit\n    EMBED_DIM = 320     # 320 for 8M model, 1280 for 650M model\n    NUM_LABELS = 2000   # Predict top 2000 most frequent GO terms\n    BATCH_SIZE = 32     # Lower if Out of Memory\n    LEARNING_RATE = 1e-3\n    EPOCHS = 5\n    \n    # Set to True to test code quickly on just 500 samples\n    DEBUG_MODE = True \n\n# ==========================================\n# 1. ONTOLOGY & LABEL PROCESSING\n# ==========================================\nclass LabelManager:\n    def __init__(self, obo_path, terms_path, num_labels=2000):\n        print(\">> Loading Gene Ontology Graph...\")\n        self.graph = obonet.read_obo(obo_path)\n        \n        # Precompute ancestors for fast lookups\n        self.term_ancestors = {\n            term: set(nx.ancestors(self.graph, term)) \n            for term in tqdm(self.graph.nodes, desc=\"Indexing Ancestors\")\n        }\n        \n        print(\">> Loading Training Terms...\")\n        self.terms_df = pd.read_csv(terms_path, sep=\"\\t\")\n        \n        # Select Top-K most frequent terms\n        term_counts = Counter(self.terms_df['term'])\n        self.selected_terms = [t[0] for t in term_counts.most_common(num_labels)]\n        self.term_to_idx = {t: i for i, t in enumerate(self.selected_terms)}\n        self.idx_to_term = {i: t for t, i in self.term_to_idx.items()}\n        print(f\"Selected {len(self.selected_terms)} most frequent GO terms.\")\n\n    def get_labels_dict(self):\n        \"\"\"Returns {protein_id: binary_vector} with hierarchy propagated\"\"\"\n        protein_groups = self.terms_df.groupby('protein_id')['term'].apply(set)\n        labels_dict = {}\n        \n        print(\">> Propagating Labels (Hierarchy Enforcement)...\")\n        # Use only subset if in Debug Mode\n        all_ids = list(protein_groups.index)\n        if Config.DEBUG_MODE: all_ids = all_ids[:500]\n            \n        for pid in tqdm(all_ids):\n            raw_terms = protein_groups[pid]\n            propagated_terms = set(raw_terms)\n            \n            # Add all parents (if Child is true, Parent is true)\n            for term in raw_terms:\n                if term in self.term_ancestors:\n                    propagated_terms.update(self.term_ancestors[term])\n            \n            # Create Binary Vector\n            vec = np.zeros(Config.NUM_LABELS, dtype=np.float32)\n            for term in propagated_terms:\n                if term in self.term_to_idx:\n                    vec[self.term_to_idx[term]] = 1.0\n            labels_dict[pid] = vec\n            \n        return labels_dict\n\n# ==========================================\n# 2. FEATURE EXTRACTION (ESM-2)\n# ==========================================\ndef extract_embeddings(fasta_path, model_name, limit=None):\n    \"\"\"\n    Runs sequences through ESM-2 and returns (ids, embeddings).\n    This mimics the 'Brain' of the operation.\n    \"\"\"\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\">> Extracting Features using {model_name} on {device}...\")\n    \n    tokenizer = AutoTokenizer.from_pretrained(model_name)\n    model = EsmModel.from_pretrained(model_name).to(device)\n    model.eval()\n    \n    seqs = []\n    ids = []\n    \n    # Parse FASTA\n    for i, record in enumerate(SeqIO.parse(fasta_path, \"fasta\")):\n        if limit and i >= limit: break\n        ids.append(str(record.id))\n        seqs.append(str(record.seq)[:Config.MAX_SEQ_LEN]) # Truncate\n        \n    # Batch Processing\n    embeddings = []\n    batch_size = 16 # Keep small for inference\n    \n    with torch.no_grad():\n        for i in tqdm(range(0, len(seqs), batch_size), desc=\"Embedding Batches\"):\n            batch_seqs = seqs[i : i + batch_size]\n            inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True)\n            inputs = {k: v.to(device) for k, v in inputs.items()}\n            \n            outputs = model(**inputs)\n            # Use the \"CLS\" token (index 0) as the full protein representation\n            batch_emb = outputs.last_hidden_state[:, 0, :].cpu().numpy()\n            embeddings.append(batch_emb)\n            \n    return ids, np.vstack(embeddings)\n\n# ==========================================\n# 3. PYTORCH DATASET & MODEL\n# ==========================================\nclass ProteinDataset(Dataset):\n    def __init__(self, embeddings, labels=None):\n        self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n        # labels can be None for Test set\n        self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n        \n    def __len__(self):\n        return len(self.embeddings)\n    \n    def __getitem__(self, idx):\n        if self.labels is not None:\n            return self.embeddings[idx], self.labels[idx]\n        return self.embeddings[idx]\n\nclass MultiLabelClassifier(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super().__init__()\n        self.layer1 = nn.Linear(input_dim, 512)\n        self.bn1 = nn.BatchNorm1d(512)\n        self.relu = nn.ReLU()\n        self.dropout = nn.Dropout(0.3)\n        self.layer2 = nn.Linear(512, num_classes)\n        # No Sigmoid here, we use BCEWithLogitsLoss for numerical stability\n        \n    def forward(self, x):\n        x = self.layer1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.dropout(x)\n        x = self.layer2(x)\n        return x\n\n# ==========================================\n# 4. MAIN EXECUTION PIPELINE\n# ==========================================\ndef main():\n    print(\"=== STARTING CAFA 6 PIPELINE ===\")\n    \n    # --- Step A: Prepare Labels ---\n    label_mgr = LabelManager(Config.OBO_FILE, Config.TRAIN_TERMS, Config.NUM_LABELS)\n    labels_map = label_mgr.get_labels_dict()\n    \n    # --- Step B: Prepare Training Data ---\n    # In a real run, set limit=None. For demo/debug, limit=500.\n    limit = 500 if Config.DEBUG_MODE else None\n    \n    train_ids, train_embeds = extract_embeddings(Config.TRAIN_FASTA, Config.MODEL_NAME, limit=limit)\n    \n    # Match embeddings to labels (Filter out sequences that don't have labels loaded)\n    X_data = []\n    y_data = []\n    \n    print(\">> Aligning Embeddings with Labels...\")\n    for pid, emb in zip(train_ids, train_embeds):\n        if pid in labels_map:\n            X_data.append(emb)\n            y_data.append(labels_map[pid])\n            \n    X_data = np.array(X_data)\n    y_data = np.array(y_data)\n    \n    # Create Dataset\n    dataset = ProteinDataset(X_data, y_data)\n    train_size = int(0.8 * len(dataset))\n    val_size = len(dataset) - train_size\n    train_ds, val_ds = random_split(dataset, [train_size, val_size])\n    \n    train_loader = DataLoader(train_ds, batch_size=Config.BATCH_SIZE, shuffle=True)\n    val_loader = DataLoader(val_ds, batch_size=Config.BATCH_SIZE)\n    \n    # --- Step C: Training the Model ---\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = MultiLabelClassifier(input_dim=Config.EMBED_DIM, num_classes=Config.NUM_LABELS).to(device)\n    \n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = optim.Adam(model.parameters(), lr=Config.LEARNING_RATE)\n    \n    print(f\"\\n>> Training for {Config.EPOCHS} epochs on {device}...\")\n    \n    for epoch in range(Config.EPOCHS):\n        model.train()\n        train_loss = 0\n        for inputs, targets in train_loader:\n            inputs, targets = inputs.to(device), targets.to(device)\n            optimizer.zero_grad()\n            outputs = model(inputs)\n            loss = criterion(outputs, targets)\n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item()\n            \n        # Validation Loop\n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for inputs, targets in val_loader:\n                inputs, targets = inputs.to(device), targets.to(device)\n                outputs = model(inputs)\n                loss = criterion(outputs, targets)\n                val_loss += loss.item()\n                \n        print(f\"Epoch {epoch+1}/{Config.EPOCHS} | Train Loss: {train_loss/len(train_loader):.4f} | Val Loss: {val_loss/len(val_loader):.4f}\")\n\n    # Save the trained brain\n    torch.save(model.state_dict(), \"cafa6_model.pth\")\n    print(\">> Model Saved!\")\n\n    # --- Step D: Inference on Test Set ---\n    print(\"\\n>> Processing Test Data...\")\n    \n    # In Debug mode, we only predict on a tiny subset of test data\n    test_limit = 200 if Config.DEBUG_MODE else None\n    \n    test_ids, test_embeds = extract_embeddings(Config.TEST_FASTA, Config.MODEL_NAME, limit=test_limit)\n    test_dataset = ProteinDataset(test_embeds, labels=None)\n    test_loader = DataLoader(test_dataset, batch_size=Config.BATCH_SIZE, shuffle=False)\n    \n    predictions = []\n    model.eval()\n    with torch.no_grad():\n        for inputs in test_loader:\n            inputs = inputs.to(device)\n            outputs = model(inputs)\n            probs = torch.sigmoid(outputs).cpu().numpy()\n            predictions.append(probs)\n            \n    predictions = np.vstack(predictions)\n    \n    # --- Step E: Submission File Generation ---\n    print(\">> Generating submission.tsv...\")\n    with open('submission.tsv', 'w') as f:\n        # No header required for CAFA\n        for i, pid in enumerate(tqdm(test_ids)):\n            # Get probabilities for this protein\n            probs = predictions[i]\n            \n            # Filter for top predictions to keep file size small (Top 50 per protein)\n            top_indices = np.argsort(probs)[-50:][::-1]\n            \n            for idx in top_indices:\n                score = probs[idx]\n                # Only write reasonable confidence scores\n                if score > 0.01:\n                    term_id = label_mgr.idx_to_term[idx]\n                    # Format: ProteinID <tab> TermID <tab> Score\n                    f.write(f\"{pid}\\t{term_id}\\t{score:.3f}\\n\")\n                    \n    print(\"=== PIPELINE COMPLETE. READY FOR SUBMISSION ===\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}