{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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":"2026-01-07T06:54:41.165977Z","iopub.execute_input":"2026-01-07T06:54:41.166361Z","iopub.status.idle":"2026-01-07T06:54:42.110058Z","shell.execute_reply.started":"2026-01-07T06:54:41.166325Z","shell.execute_reply":"2026-01-07T06:54:42.109279Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Install necessary libraries\n!pip install -q biopython transformers torch numpy pandas tqdm scikit-learn","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:54:42.112059Z","iopub.execute_input":"2026-01-07T06:54:42.112482Z","iopub.status.idle":"2026-01-07T06:54:48.404597Z","shell.execute_reply.started":"2026-01-07T06:54:42.112455Z","shell.execute_reply":"2026-01-07T06:54:48.403890Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom Bio import SeqIO\nfrom tqdm import tqdm\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModel\nimport torch.nn as nn\nimport matplotlib.pyplot as plt\n\n# Configuration\nCONFIG = {\n    # CHANGE 1: Use the 35M model (The \"Goldilocks\" model)\n    \"model_name\": \"facebook/esm2_t12_35M_UR50D\", \n    \n    # CHANGE 2: Increase Batch Size (35M fits easily)\n    \"batch_size\": 8,          \n    \n    # CHANGE 3: Train longer!\n    \"epochs\": 3,\n    \n    \"max_len\": 512,\n    \"learning_rate\": 5e-4, # Slightly higher LR for this size\n    \"num_labels\": 2000,\n    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n    \"paths\": {\n        \"train_seq\": \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\",\n        \"train_terms\": \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n        \"test_seq\": \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\",\n        \"sample_sub\": \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\"\n    }\n}\n\nprint(f\"Using device: {CONFIG['device']}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:54:48.405835Z","iopub.execute_input":"2026-01-07T06:54:48.406169Z","iopub.status.idle":"2026-01-07T06:55:00.048125Z","shell.execute_reply.started":"2026-01-07T06:54:48.406124Z","shell.execute_reply":"2026-01-07T06:55:00.047217Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_fasta(path):\n    \"\"\"Parses a FASTA file into a dictionary: {id: sequence}\"\"\"\n    sequences = {}\n    print(f\"Loading sequences from {path}...\")\n    # Biopython's SeqIO makes this easy\n    for record in SeqIO.parse(path, \"fasta\"):\n        # CAFA headers look like \"sp|P12345|NAME\". We just want \"P12345\"\n        # Adjust parsing logic if headers are different, but usually ID is the 2nd part\n        header_parts = record.id.split(\"|\")\n        if len(header_parts) >= 2:\n            pid = header_parts[1]\n        else:\n            pid = record.id # Fallback\n        sequences[pid] = str(record.seq)\n    print(f\"Loaded {len(sequences)} sequences.\")\n    return sequences\n\n# Load Train Data\ntrain_seqs = load_fasta(CONFIG['paths']['train_seq'])\ntrain_terms = pd.read_csv(CONFIG['paths']['train_terms'], sep=\"\\t\")\n\nprint(\"Train Terms Head:\")\nprint(train_terms.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:55:00.049297Z","iopub.execute_input":"2026-01-07T06:55:00.049888Z","iopub.status.idle":"2026-01-07T06:55:01.258287Z","shell.execute_reply.started":"2026-01-07T06:55:00.049854Z","shell.execute_reply":"2026-01-07T06:55:01.257468Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. Count term frequencies\nterm_counts = train_terms['term'].value_counts()\n\n# 2. Select top N terms\ntarget_terms = term_counts.head(CONFIG['num_labels']).index.tolist()\nterm_to_idx = {term: i for i, term in enumerate(target_terms)}\nidx_to_term = {i: term for i, term in enumerate(target_terms)}\n\nprint(f\"Selected top {len(target_terms)} GO terms for prediction.\")\n\n# 3. Create a valid training dataframe\n# We only want proteins that have sequences AND annotations in our target list\nfiltered_terms = train_terms[train_terms['term'].isin(target_terms)]\nprotein_to_terms = filtered_terms.groupby('EntryID')['term'].apply(list).to_dict()\n\n# Filter sequences to those that have labels\ntrain_ids = [pid for pid in train_seqs.keys() if pid in protein_to_terms]\nprint(f\"Training on {len(train_ids)} proteins that have labels in our top-{CONFIG['num_labels']} set.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:55:01.259318Z","iopub.execute_input":"2026-01-07T06:55:01.259769Z","iopub.status.idle":"2026-01-07T06:55:02.631045Z","shell.execute_reply.started":"2026-01-07T06:55:01.259741Z","shell.execute_reply":"2026-01-07T06:55:02.630340Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Download the Gene Ontology structure (OBO file)\n!wget http://purl.obolibrary.org/obo/go/go-basic.obo -O go-basic.obo","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:55:02.631937Z","iopub.execute_input":"2026-01-07T06:55:02.632189Z","iopub.status.idle":"2026-01-07T06:55:03.336836Z","shell.execute_reply.started":"2026-01-07T06:55:02.632165Z","shell.execute_reply":"2026-01-07T06:55:03.336018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ProteinGO_Dataset(Dataset):\n    def __init__(self, protein_ids, sequences, protein_to_terms, tokenizer, term_to_idx, max_len):\n        self.protein_ids = protein_ids\n        self.sequences = sequences\n        self.protein_to_terms = protein_to_terms\n        self.tokenizer = tokenizer\n        self.term_to_idx = term_to_idx\n        self.max_len = max_len\n        self.num_classes = len(term_to_idx)\n\n    def __len__(self):\n        return len(self.protein_ids)\n\n    def __getitem__(self, idx):\n        pid = self.protein_ids[idx]\n        seq = self.sequences[pid]\n        \n        # Tokenize sequence\n        encoding = self.tokenizer(\n            seq, \n            truncation=True, \n            padding='max_length', \n            max_length=self.max_len, \n            return_tensors='pt'\n        )\n        \n        # Create Multi-Hot Label Vector\n        labels = torch.zeros(self.num_classes)\n        if self.protein_to_terms: # If training mode\n            annotated_terms = self.protein_to_terms.get(pid, [])\n            for term in annotated_terms:\n                if term in self.term_to_idx:\n                    labels[self.term_to_idx[term]] = 1.0\n        \n        return {\n            'input_ids': encoding['input_ids'].squeeze(0),\n            'attention_mask': encoding['attention_mask'].squeeze(0),\n            'labels': labels,\n            'pid': pid\n        }\n\n# Initialize Tokenizer\ntokenizer = AutoTokenizer.from_pretrained(CONFIG['model_name'])\n\n# Create Train/Val split (Simple 80/20)\nnp.random.seed(42)\nnp.random.shuffle(train_ids)\nsplit_idx = int(len(train_ids) * 0.8)\ntrain_split = train_ids[:split_idx]\nval_split = train_ids[split_idx:]\n\ntrain_dataset = ProteinGO_Dataset(train_split, train_seqs, protein_to_terms, tokenizer, term_to_idx, CONFIG['max_len'])\nval_dataset = ProteinGO_Dataset(val_split, train_seqs, protein_to_terms, tokenizer, term_to_idx, CONFIG['max_len'])\n\ntrain_loader = DataLoader(train_dataset, batch_size=CONFIG['batch_size'], shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=CONFIG['batch_size'], shuffle=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:55:03.339776Z","iopub.execute_input":"2026-01-07T06:55:03.340282Z","iopub.status.idle":"2026-01-07T06:55:04.199152Z","shell.execute_reply.started":"2026-01-07T06:55:03.340247Z","shell.execute_reply":"2026-01-07T06:55:04.198491Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- UPDATED MODEL (NO SIGMOID) ---\nclass ESM_GO_Predictor(nn.Module):\n    def __init__(self, model_name, num_classes):\n        super().__init__()\n        # Load ESM model\n        self.esm = AutoModel.from_pretrained(model_name)\n        \n        # Freezing weights (Optional: Unfreeze for better accuracy if you have time)\n        for param in self.esm.parameters():\n            param.requires_grad = False\n            \n        self.embedding_dim = self.esm.config.hidden_size\n        \n        # Classifier Head\n        self.classifier = nn.Sequential(\n            nn.Linear(self.embedding_dim, 512),\n            nn.ReLU(),\n            nn.Dropout(0.2),\n            nn.Linear(512, num_classes) \n            # REMOVED nn.Sigmoid() -> We rely on BCEWithLogitsLoss now\n        )\n\n    def forward(self, input_ids, attention_mask):\n        outputs = self.esm(input_ids=input_ids, attention_mask=attention_mask)\n        embeddings = outputs.last_hidden_state\n        \n        # Masked Mean Pooling\n        mask = attention_mask.unsqueeze(-1).expand(embeddings.size()).float()\n        masked_embeddings = embeddings * mask\n        sum_embeddings = torch.sum(masked_embeddings, 1)\n        sum_mask = torch.clamp(mask.sum(1), min=1e-9)\n        mean_embeddings = sum_embeddings / sum_mask\n        \n        return self.classifier(mean_embeddings)\n\n# Re-initialize the model with this new structure\nmodel = ESM_GO_Predictor(CONFIG['model_name'], CONFIG['num_labels'])\nmodel.to(CONFIG['device'])\nprint(\"✅ Model re-initialized without Sigmoid layer.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:55:04.200187Z","iopub.execute_input":"2026-01-07T06:55:04.200491Z","iopub.status.idle":"2026-01-07T06:55:23.249959Z","shell.execute_reply.started":"2026-01-07T06:55:04.200465Z","shell.execute_reply":"2026-01-07T06:55:23.249128Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- STABLE TRAINING LOOP (35M Model) ---\nfrom torch.amp import autocast, GradScaler\nimport torch.nn as nn\n\n# 1. Optimizer & Loss\noptimizer = torch.optim.Adam(model.parameters(), lr=CONFIG['learning_rate'])\ncriterion = nn.BCEWithLogitsLoss() # The stable loss function\nscaler = GradScaler('cuda')\n\n# 2. Training Function\ndef train_epoch_amp(model, loader, optimizer, criterion):\n    model.train()\n    total_loss = 0\n    pbar = tqdm(loader, desc=\"Training (35M Model)\")\n    \n    for batch in pbar:\n        input_ids = batch['input_ids'].to(CONFIG['device'])\n        attention_mask = batch['attention_mask'].to(CONFIG['device'])\n        labels = batch['labels'].to(CONFIG['device'])\n        \n        optimizer.zero_grad()\n        \n        # Mixed Precision\n        with autocast('cuda'):\n            outputs = model(input_ids, attention_mask)\n            loss = criterion(outputs, labels)\n        \n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        \n        total_loss += loss.item()\n        pbar.set_postfix({'loss': f\"{loss.item():.4f}\"})\n        \n    return total_loss / len(loader)\n\n# 3. Run for 3 Epochs\nprint(f\"🚀 Starting training for {CONFIG['epochs']} epochs...\")\n\nfor epoch in range(CONFIG['epochs']):\n    train_loss = train_epoch_amp(model, train_loader, optimizer, criterion)\n    print(f\"Epoch {epoch+1} Completed | Train Loss: {train_loss:.4f}\")\n    \n    # Save each epoch\n    torch.save(model.state_dict(), f\"esm2_t12_epoch_{epoch+1}.pth\")\n\nprint(\"✅ Training finished.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T06:55:23.251025Z","iopub.execute_input":"2026-01-07T06:55:23.251598Z","iopub.status.idle":"2026-01-07T07:23:42.950638Z","shell.execute_reply.started":"2026-01-07T06:55:23.251571Z","shell.execute_reply":"2026-01-07T07:23:42.949972Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- HIGH-SPEED INFERENCE (Dynamic Padding + Sorting) ---\nfrom torch.nn.utils.rnn import pad_sequence\n\n# 1. Load Test Data\ntest_seqs = load_fasta(CONFIG['paths']['test_seq'])\ntest_ids = list(test_seqs.keys())\n\n# 2. SORT sequences by length\n# This groups short proteins together so they process super fast\nprint(\"Sorting test data by length for speed...\")\ntest_ids.sort(key=lambda x: len(test_seqs[x]))\n\n# 3. Dynamic Collate Function (The Secret Sauce)\n# Instead of padding to 512, we pad to the max length in the BATCH\ndef dynamic_collate_fn(batch):\n    # Separate inputs\n    input_ids = [item['input_ids'] for item in batch]\n    attention_mask = [item['attention_mask'] for item in batch]\n    pids = [item['pid'] for item in batch]\n    \n    # Pad dynamically using PyTorch util\n    # padding_value=1 is usually the <pad> token for ESM (verify if 0 or 1, esm tokenizer usually uses 1 for pad)\n    # We can check tokenizer.pad_token_id. usually 1.\n    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)\n    attention_mask_padded = pad_sequence(attention_mask, batch_first=True, padding_value=0)\n    \n    return {\n        'input_ids': input_ids_padded,\n        'attention_mask': attention_mask_padded,\n        'pid': pids\n    }\n\n# 4. Modified Dataset (Returns unpadded tensors)\nclass ProteinGO_Inference_Dataset(Dataset):\n    def __init__(self, protein_ids, sequences, tokenizer, max_len):\n        self.protein_ids = protein_ids\n        self.sequences = sequences\n        self.tokenizer = tokenizer\n        self.max_len = max_len\n\n    def __len__(self):\n        return len(self.protein_ids)\n\n    def __getitem__(self, idx):\n        pid = self.protein_ids[idx]\n        seq = self.sequences[pid]\n        \n        # Tokenize WITHOUT padding first\n        encoding = self.tokenizer(\n            seq, \n            truncation=True, \n            max_length=self.max_len, \n            return_tensors='pt'\n        )\n        \n        # Remove batch dim added by tokenizer\n        return {\n            'input_ids': encoding['input_ids'].squeeze(0),\n            'attention_mask': encoding['attention_mask'].squeeze(0),\n            'pid': pid\n        }\n\n# 5. Create Optimized Loader\ntest_dataset_opt = ProteinGO_Inference_Dataset(test_ids, test_seqs, tokenizer, CONFIG['max_len'])\n\n# Increase batch size even more since short batches will be tiny\ninference_batch = 32 \n\ntest_loader = DataLoader(\n    test_dataset_opt, \n    batch_size=inference_batch, \n    shuffle=False, \n    collate_fn=dynamic_collate_fn, # Use the smart padder\n    num_workers=2,\n    pin_memory=True\n)\n\n# 6. Run Inference with Autocast\nmodel.eval()\npredictions = []\n\nprint(f\"🚀 Running Speed-Optimized Inference (Batch {inference_batch})...\")\n\nwith torch.no_grad():\n    for batch in tqdm(test_loader, desc=\"Inference\"):\n        input_ids = batch['input_ids'].to(CONFIG['device'])\n        attention_mask = batch['attention_mask'].to(CONFIG['device'])\n        pids = batch['pid']\n        \n        # Use Mixed Precision for Inference too!\n        with torch.amp.autocast('cuda'):\n            outputs = model(input_ids, attention_mask)\n            # Add Sigmoid manually (since we removed it from model)\n            probs = torch.sigmoid(outputs)\n        \n        # Move to CPU\n        probs = probs.float().cpu().numpy()\n        \n        for i, pid in enumerate(pids):\n            prot_probs = probs[i]\n            top_indices = np.where(prot_probs > 0.01)[0] \n            for idx in top_indices:\n                predictions.append([pid, idx_to_term[idx], f\"{prot_probs[idx]:.3f}\"])\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T07:23:42.951889Z","iopub.execute_input":"2026-01-07T07:23:42.952174Z","iopub.status.idle":"2026-01-07T07:40:13.620416Z","shell.execute_reply.started":"2026-01-07T07:23:42.952148Z","shell.execute_reply":"2026-01-07T07:40:13.619661Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create DataFrame\nsubmission_df = pd.DataFrame(predictions, columns=['EntryID', 'term', 'probability'])\n\n# Save\nsubmission_df.to_csv(\"submission.tsv\", sep=\"\\t\", index=False, header=False) \n# Note: Check competition rules if header is required. Usually CAFA is headerless or specific.\n# Based on previous CAFA, typically no header or simple one. We stick to standard formats.\n\nprint(\"✅ Submission file 'submission.tsv' created successfully!\")\nprint(submission_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T07:40:13.621663Z","iopub.execute_input":"2026-01-07T07:40:13.622236Z","iopub.status.idle":"2026-01-07T07:40:27.511402Z","shell.execute_reply.started":"2026-01-07T07:40:13.622207Z","shell.execute_reply":"2026-01-07T07:40:27.510726Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# SUPER LEARNER PHASE: MERGE & BOOST (CORRECTED)\n# ==========================================\nimport os\nimport pandas as pd\nimport numpy as np\n\n# --- 1. SETUP DIAMOND ---\nprint(\"💎 Installing Diamond...\")\nif not os.path.exists(\"diamond\"):\n    !wget -q http://github.com/bbuchfink/diamond/releases/download/v2.1.8/diamond-linux64.tar.gz\n    !tar xzf diamond-linux64.tar.gz\n    !rm diamond-linux64.tar.gz\n    !chmod +x diamond\n\n# Paths (Double check these match your Kaggle input!)\nTRAIN_FASTA = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\nTRAIN_TERMS = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\nTEST_FASTA  = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n\n# --- 2. RUN DIAMOND SEARCH ---\nprint(\"💎 Building Diamond Database...\")\n# Check if DB exists to save time, else build\nif not os.path.exists(\"train_db.dmnd\"):\n    !./diamond makedb --in {TRAIN_FASTA} -d train_db --quiet\n\nprint(\"💎 Running Alignment (Test vs Train)...\")\n# Run Diamond. Output columns: qseqid (TestID), sseqid (TrainID), pident (Identity)\n!./diamond blastp -d train_db -q {TEST_FASTA} -o diamond_matches.tsv -k 10 --quiet --outfmt 6 qseqid sseqid pident\n\n# --- 3. PROCESS DIAMOND RESULTS WITH ID CLEANING ---\nprint(\"💎 Processing Alignment Results...\")\n\n# Load matches\n# Columns: TestID, TrainID, Identity\ntry:\n    matches = pd.read_csv(\"diamond_matches.tsv\", sep=\"\\t\", header=None, names=[\"TestID\", \"TrainID\", \"Identity\"])\nexcept pd.errors.EmptyDataError:\n    print(\"⚠️ WARNING: Diamond found NO matches. Check your paths!\")\n    matches = pd.DataFrame(columns=[\"TestID\", \"TrainID\", \"Identity\"])\n\n# --- CRITICAL FIX: CLEAN IDs ---\ndef clean_id(x):\n    # Converts \"sp|P12345|NAME\" -> \"P12345\"\n    # Converts \"P12345\" -> \"P12345\"\n    x = str(x)\n    if \"|\" in x:\n        parts = x.split(\"|\")\n        # Usually the ID is the 2nd part (index 1) in Uniprot headers\n        if len(parts) >= 2:\n            return parts[1]\n    return x\n\n# Apply cleaning to both Train and Test IDs from Diamond output\nmatches['TrainID'] = matches['TrainID'].apply(clean_id)\nmatches['TestID'] = matches['TestID'].apply(clean_id)\n\nprint(f\"   -> Loaded {len(matches)} raw matches after cleaning IDs.\")\n\n# Load Train Terms\ntrain_terms_df = pd.read_csv(TRAIN_TERMS, sep=\"\\t\")\ntrain_annotations = train_terms_df.groupby(\"EntryID\")[\"term\"].apply(list).to_dict()\n\ndiamond_preds = []\nfor _, row in matches.iterrows():\n    test_id = row[\"TestID\"]\n    train_id = row[\"TrainID\"]\n    identity = row[\"Identity\"] / 100.0\n    \n    # Now the lookup should work!\n    if train_id in train_annotations:\n        annotated_terms = train_annotations[train_id]\n        for term in annotated_terms:\n            diamond_preds.append([test_id, term, identity])\n\n# Convert to DataFrame\ndf_diamond = pd.DataFrame(diamond_preds, columns=[\"EntryID\", \"term\", \"score\"])\nif not df_diamond.empty:\n    df_diamond = df_diamond.groupby([\"EntryID\", \"term\"])[\"score\"].max().reset_index()\nprint(f\"💎 Diamond found predictions for {df_diamond['EntryID'].nunique()} proteins.\")\n\n# --- 4. MERGE WITH DEEP LEARNING ---\nprint(\"🤝 Merging with Deep Learning Predictions...\")\n\n# Load your ESM predictions\nif os.path.exists(\"submission.tsv\"):\n    df_dl = pd.read_csv(\"submission.tsv\", sep=\"\\t\", header=None, names=[\"EntryID\", \"term\", \"score\"])\nelse:\n    print(\"⚠️ submission.tsv not found! Running with only Diamond (risky).\")\n    df_dl = pd.DataFrame(columns=[\"EntryID\", \"term\", \"score\"])\n\n# Combine\ncombined_df = pd.concat([df_dl, df_diamond])\n\n# Max Pooling\nfinal_df = combined_df.groupby([\"EntryID\", \"term\"])[\"score\"].max().reset_index()\n\n# --- 5. SAVE ---\nprint(\"📝 saving final_submission.tsv...\")\nfinal_df.to_csv(\"final_submission.tsv\", sep=\"\\t\", index=False, header=False)\n\nprint(f\"✅ DONE! Check the output file.\")\nprint(final_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T07:40:27.512318Z","iopub.execute_input":"2026-01-07T07:40:27.512536Z","iopub.status.idle":"2026-01-07T07:45:13.247679Z","shell.execute_reply.started":"2026-01-07T07:40:27.512514Z","shell.execute_reply":"2026-01-07T07:45:13.246956Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install obonet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T07:45:13.248740Z","iopub.execute_input":"2026-01-07T07:45:13.249062Z","iopub.status.idle":"2026-01-07T07:45:16.605141Z","shell.execute_reply.started":"2026-01-07T07:45:13.249034Z","shell.execute_reply":"2026-01-07T07:45:16.604416Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import csv\nimport obonet\nimport networkx as nx\nfrom itertools import groupby\nfrom tqdm import tqdm\nimport os\nimport gzip # <--- The secret weapon\n\n# --- 1. CLEANUP (Critical) ---\n# Delete the failed file to free up space immediately\nif os.path.exists(\"submission_propagated.tsv\"):\n    os.remove(\"submission_propagated.tsv\")\n    print(\"🗑️ Deleted broken/partial file to clear space.\")\n\n# --- 2. SETUP ---\nif not os.path.exists(\"final_submission.tsv\"):\n    print(\"❌ ERROR: Input file missing. Please check.\")\nelse:\n    print(\"✅ Found input. Starting Compressed & Limited Propagation...\")\n\n    # Load Graph\n    graph = obonet.read_obo(\"go-basic.obo\")\n    parents_map = {node: list(graph.predecessors(node)) for node in graph.nodes()}\n\n    input_file = \"final_submission.tsv\"\n    # We write to a GZIP file now to save disk space\n    output_file = \"submission_propagated.tsv.gz\" \n\n    def read_tsv_generator(filename):\n        with open(filename, 'r') as f:\n            reader = csv.reader(f, delimiter='\\t')\n            for row in reader:\n                yield row\n\n    print(\"🚀 Starting Smart Propagation (Compressed + Top 500 Limit)...\")\n    \n    # Open with gzip (mode 'wt' means Write Text)\n    with gzip.open(output_file, 'wt', newline='') as f_out:\n        writer = csv.writer(f_out, delimiter='\\t')\n        data_stream = read_tsv_generator(input_file)\n        \n        for pid, group in tqdm(groupby(data_stream, key=lambda x: x[0]), desc=\"Propagating\"):\n            terms_dict = {}\n            for row in group:\n                try:\n                    terms_dict[row[1]] = float(row[2])\n                except ValueError: continue\n\n            # Propagate Scores\n            current_terms = list(terms_dict.keys())\n            for _ in range(3): \n                for term in current_terms:\n                    if term in parents_map:\n                        current_score = terms_dict[term]\n                        for parent in parents_map[term]:\n                            parent_score = terms_dict.get(parent, 0.0)\n                            if current_score > parent_score:\n                                terms_dict[parent] = current_score\n                                current_terms.append(parent)\n            \n            # --- SAFETY LIMIT: Sort & Keep Top 500 ---\n            # 1. Convert dict to list of tuples: [('GO:001', 0.9), ('GO:002', 0.8)...]\n            sorted_terms = sorted(terms_dict.items(), key=lambda x: x[1], reverse=True)\n            \n            # 2. Keep only Top 500 (Most valuable predictions)\n            top_terms = sorted_terms[:500] \n\n            # 3. Write\n            for term, score in top_terms:\n                if score >= 0.01: # Keep the 0.01 threshold\n                    writer.writerow([pid, term, f\"{score:.3f}\"])\n\n    print(\"✅ SUCCESS! Created 'submission_propagated.tsv.gz'\")\n    print(\"👉 ACTION: Download this file, UNZIP it on your computer, and submit the extracted file.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T07:45:16.606484Z","iopub.execute_input":"2026-01-07T07:45:16.606770Z","iopub.status.idle":"2026-01-07T08:53:39.086689Z","shell.execute_reply.started":"2026-01-07T07:45:16.606739Z","shell.execute_reply":"2026-01-07T08:53:39.085996Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}