{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.5"},"kaggle":{"accelerator":"gpu","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":"markdown","source":"# CAFA 6 Protein Function Prediction: Robust Ensemble Solution\n### Author: Antigravity\n\nThis notebook implements a high-performance ensemble for the CAFA 6 competition.\nIt combines:\n1.  **Similarity Search (Diamond)**: High precision for homologs.\n2.  **ResNet (Deep Learning)**: Deep residual network on ESM-2 embeddings.\n3.  **PyTorch MLP**: Simple coverage model.\n4.  **Naive Baseline**: Prior term probabilities (Critical for F-max stability).\n\n## Setup\nUncomment the lines below to install dependencies if running on a fresh environment.","metadata":{}},{"cell_type":"code","source":"# !conda install -c bioconda diamond -y\n# !pip install torch transformers pandas numpy tqdm scipy biopython","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport subprocess\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom collections import defaultdict\nfrom Bio import SeqIO\nfrom transformers import AutoTokenizer, AutoModel\nfrom sklearn.model_selection import train_test_split\n\n# Detect Device\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Running on: {DEVICE}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- CONFIGURATION (Adjust for Kaggle) ---\n# Inputs\nINPUT_DIR = '../input/cafa-6-protein-function-prediction' \n# If you uploaded embeddings/models as a dataset, point to them:\nDATASET_DIR = '../input/cafa6-embeddings' # Example path\n\nTRAIN_FASTA = f\"{INPUT_DIR}/Train/train_sequences.fasta\"\nTRAIN_TERMS = f\"{INPUT_DIR}/Train/train_terms.tsv\"\nTEST_FASTA = f\"{INPUT_DIR}/Test/testsuperset.fasta\"\n\n# Outputs\nOUTPUT_DIR = './' # /kaggle/working/\nEMBED_DIR = './embeddings'\nOS.makedirs(EMBED_DIR, exist_ok=True)\n\n# Local Fallback (for testing)\nif not os.path.exists(INPUT_DIR):\n    INPUT_DIR = '/home/hiroshige/anti/cafa6'\n    TRAIN_FASTA = f\"{INPUT_DIR}/Train/train_sequences.fasta\"\n    TRAIN_TERMS = f\"{INPUT_DIR}/Train/train_terms.tsv\"\n    TEST_FASTA = f\"{INPUT_DIR}/Test/testsuperset.fasta\"\n    OUTPUT_DIR = INPUT_DIR\n    EMBED_DIR = INPUT_DIR\n\nprint(f\"Input Dir: {INPUT_DIR}\")\nprint(f\"Output Dir: {OUTPUT_DIR}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- ESM-2 EMBEDDING GENERATION ---\nMODEL_NAME = \"facebook/esm2_t6_8M_UR50D\"\nMAX_LEN = 1024\nBATCH_SIZE = 32\n\ndef generate_embeddings(fasta_file, output_prefix):\n    ids_file = f\"{output_prefix}_ids.npy\"\n    embed_file = f\"{output_prefix}_embeddings.npy\"\n    \n    if os.path.exists(embed_file) and os.path.exists(ids_file):\n        print(f\"Embeddings found for {fasta_file}: {embed_file}\")\n        return np.load(ids_file, allow_pickle=True), np.load(embed_file, allow_pickle=True)\n        \n    print(f\"Generating embeddings for {fasta_file}...\")\n    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n    model = AutoModel.from_pretrained(MODEL_NAME).to(DEVICE)\n    model.eval()\n\n    sequences = []\n    ids = []\n    for record in SeqIO.parse(fasta_file, \"fasta\"):\n        sequences.append(str(record.seq)[:MAX_LEN])\n        ids.append(record.id.split('|')[1] if '|' in record.id else record.id)\n    \n    all_embeddings = []\n    with torch.no_grad():\n        for i in tqdm(range(0, len(sequences), BATCH_SIZE)):\n            batch_seqs = sequences[i:i+BATCH_SIZE]\n            inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=MAX_LEN)\n            inputs = {k: v.to(DEVICE) for k, v in inputs.items()}\n            outputs = model(**inputs)\n            \n            # Mean pooling with mask\n            attention_mask = inputs['attention_mask'].unsqueeze(-1).float()\n            masked = outputs.last_hidden_state * attention_mask\n            pooled = masked.sum(dim=1) / attention_mask.sum(dim=1)\n            all_embeddings.append(pooled.cpu().numpy().astype(np.float32))\n            \n    final_embeddings = np.concatenate(all_embeddings, axis=0)\n    \n    # Save cache\n    np.save(ids_file, np.array(ids))\n    np.save(embed_file, final_embeddings)\n    return np.array(ids), final_embeddings","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- NAIVE BASELINE ---\ndef run_naive(test_ids, output_file):\n    print(\"Running Naive Baseline...\")\n    df = pd.read_csv(TRAIN_TERMS, sep='\\t')\n    total_proteins = df['EntryID'].nunique()\n    term_counts = df['term'].value_counts()\n    term_probs = term_counts / total_proteins\n    \n    # Filter 0.5% threshold\n    top_terms = term_probs[term_probs > 0.005]\n    print(f\"Naive: Found {len(top_terms)} terms > 0.5%\")\n    \n    term_strings = [f\"\\t{term}\\t{score:.5f}\\n\" for term, score in top_terms.items()]\n    \n    with open(output_file, 'w') as f:\n        for pid in tqdm(test_ids):\n            for t_str in term_strings:\n                f.write(f\"{pid}{t_str}\")\n    print(f\"Naive Saved: {output_file}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- RESNET MODEL ---\nclass ResBlock(nn.Module):\n    def __init__(self, dim, dropout=0.3):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.BatchNorm1d(dim),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n    def forward(self, x):\n        return x + self.block(x)\n\nclass ResNet(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super().__init__()\n        self.entry = nn.Sequential(\n            nn.Linear(input_dim, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.2)\n        )\n        self.res_blocks = nn.Sequential(ResBlock(512), ResBlock(512))\n        self.exit = nn.Linear(512, num_classes)\n        \n    def forward(self, x):\n        x = self.entry(x)\n        x = self.res_blocks(x)\n        return self.exit(x)\n\ndef train_predict_resnet(train_ids, X_train, test_ids, X_test, output_file, top_k=500):\n    print(\"Preparing ResNet...\")\n    # Prepare Labels\n    df = pd.read_csv(TRAIN_TERMS, sep='\\t')\n    top_terms = df['term'].value_counts().head(top_k).index.tolist()\n    term_to_idx = {term: i for i, term in enumerate(top_terms)}\n    \n    df = df[df['EntryID'].isin(set(train_ids))]\n    id_groups = df.groupby('EntryID')['term'].apply(list).to_dict()\n    \n    y_train = np.zeros((len(train_ids), top_k), dtype=int)\n    for i, pid in enumerate(tqdm(train_ids, desc=\"Encoding Labels\")):\n        if pid in id_groups:\n            for term in id_groups[pid]:\n                if term in term_to_idx:\n                    y_train[i, term_to_idx[term]] = 1\n                    \n    # Training\n    X_train_t, X_val_t, y_train_t, y_val_t = train_test_split(X_train, y_train, test_size=0.1)\n    train_loader = DataLoader(list(zip(torch.FloatTensor(X_train_t), torch.FloatTensor(y_train_t))), batch_size=256, shuffle=True)\n    val_loader = DataLoader(list(zip(torch.FloatTensor(X_val_t), torch.FloatTensor(y_val_t))), batch_size=256)\n    \n    model = ResNet(X_train.shape[1], top_k).to(DEVICE)\n    optimizer = optim.AdamW(model.parameters(), lr=1e-3)\n    criterion = nn.BCEWithLogitsLoss()\n    \n    print(\"Training...\")\n    for epoch in range(15): # 15 epochs for Kaggle efficiency\n        model.train()\n        for bx, by in train_loader:\n            optimizer.zero_grad()\n            loss = criterion(model(bx.to(DEVICE)), by.to(DEVICE))\n            loss.backward()\n            optimizer.step()\n        \n        # Simple Validation Log\n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for bx, by in val_loader:\n                val_loss += criterion(model(bx.to(DEVICE)), by.to(DEVICE)).item()\n        print(f\"Epoch {epoch+1}: Val Loss {val_loss/len(val_loader):.4f}\")\n        \n    # Inference\n    print(\"Predicting...\")\n    test_loader = DataLoader(torch.FloatTensor(X_test), batch_size=512)\n    probs = []\n    with torch.no_grad():\n        for bx in tqdm(test_loader):\n            probs.append(torch.sigmoid(model(bx.to(DEVICE))).cpu().numpy())\n    probs = np.concatenate(probs)\n    \n    with open(output_file, 'w') as f:\n        for i, pid in enumerate(tqdm(test_ids)):\n            p = probs[i]\n            indices = np.where(p > 0.01)[0]\n            for idx in indices:\n                f.write(f\"{pid}\\t{top_terms[idx]}\\t{p[idx]:.3f}\\n\")\n    print(f\"ResNet Saved: {output_file}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- SIMILARITY SEARCH (DIAMOND) ---\ndef run_similarity(output_file):\n    print(\"Checking for Diamond...\")\n    diamond_cmd = None\n    try:\n        # Check local path first (if installed via conda)\n        import shutil\n        if shutil.which('diamond'):\n            diamond_cmd = 'diamond'\n        # Check local folder (if uploaded)\n        elif os.path.exists('./diamond'):\n            diamond_cmd = './diamond'\n            subprocess.run(['chmod', '+x', diamond_cmd])\n            \n        if not diamond_cmd:\n            raise FileNotFoundError(\"Diamond binary not found.\")\n            \n        print(f\"Diamond found: {diamond_cmd}\")\n        \n    except Exception as e:\n        print(f\"Diamond check failed: {e}. Skipping Similarity Search.\")\n        return False\n        \n    db_file = f\"{EMBED_DIR}/train.dmnd\"\n    matches_file = f\"{EMBED_DIR}/matches.tsv\"\n    \n    try:\n        if not os.path.exists(db_file):\n            print(\"Building DB...\")\n            subprocess.run([diamond_cmd, 'makedb', '--in', TRAIN_FASTA, '-d', db_file], check=True)\n            \n        if not os.path.exists(matches_file):\n            print(\"BlastP Search...\")\n            subprocess.run([diamond_cmd, 'blastp', '-d', db_file, '-q', TEST_FASTA, '-o', matches_file, \n                            '-k', '5', '--outfmt', '6', '--sensitive', '-b', '8.0'], check=True)\n                            \n        print(\"Processing Matches...\")\n        if not os.path.exists(matches_file):\n             print(\"Matches file not generated.\")\n             return False\n             \n        matches = pd.read_csv(matches_file, sep='\\t', names=['qseqid', 'sseqid', 'pident', 'length', 'mismatch', 'gapopen', 'qstart', 'qend', 'sstart', 'send', 'evalue', 'bitscore'])\n        matches['sseqid'] = matches['sseqid'].apply(lambda x: x.split('|')[1] if '|' in x else x)\n        \n        terms_df = pd.read_csv(TRAIN_TERMS, sep='\\t')\n        term_map = terms_df.groupby('EntryID')['term'].apply(list).to_dict()\n        \n        with open(output_file, 'w') as f:\n            for qid, group in tqdm(matches.groupby('qseqid')):\n                seen = {}\n                for _, row in group.iterrows():\n                     sid = row['sseqid']\n                     score = row['pident'] / 100.0\n                     if sid in term_map:\n                         for term in term_map[sid]:\n                             seen[term] = max(seen.get(term, 0), score)\n                for term, score in seen.items():\n                    f.write(f\"{qid}\\t{term}\\t{score:.3f}\\n\")\n        print(f\"Similarity Saved: {output_file}\")\n        return True\n        \n    except Exception as e:\n        print(f\"Error during Similarity Search: {e}\")\n        return False","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- ENSEMBLE ---\ndef load_preds(path):\n    d = defaultdict(dict)\n    if not os.path.exists(path): \n        print(f\"Skipping missing file: {path}\")\n        return d\n    try:\n        with open(path, 'r') as f:\n            for line in tqdm(f, desc=f\"Loading {os.path.basename(path)}\"):\n                parts = line.strip().split('\\t')\n                if len(parts) == 3:\n                    d[parts[0]][parts[1]] = float(parts[2])\n    except Exception as e:\n        print(f\"Error loading {path}: {e}\")\n    return d\n\ndef run_ensemble(out_naive, out_resnet, out_sim, final_out):\n    print(\"Loading predictions...\")\n    p_naive = load_preds(out_naive)\n    p_resnet = load_preds(out_resnet)\n    p_sim = load_preds(out_sim)\n    \n    all_ids = set(p_naive.keys()) | set(p_resnet.keys()) | set(p_sim.keys())\n    \n    if not all_ids:\n        print(\"Error: No predictions found from any model!\")\n        return\n\n    print(f\"Ensembling {len(all_ids)} proteins...\")\n    \n    try:\n        with open(final_out, 'w') as f:\n            for pid in tqdm(all_ids):\n                terms = set(p_naive.get(pid, {}).keys()) | set(p_resnet.get(pid, {}).keys()) | set(p_sim.get(pid, {}).keys())\n                for term in terms:\n                    s_naive = p_naive.get(pid, {}).get(term, 0.0)\n                    s_res = p_resnet.get(pid, {}).get(term, 0.0)\n                    s_sim = p_sim.get(pid, {}).get(term, 0.0)\n                    \n                    # Robust Weighting\n                    if not p_sim:\n                         score = s_res * 0.7 + s_naive * 0.3\n                    else:\n                         score = s_sim * 0.4 + s_res * 0.4 + s_naive * 0.2\n                         \n                    if score > 0.001:\n                        f.write(f\"{pid}\\t{term}\\t{score:.3f}\\n\")\n        print(f\"Final Submission Saved: {final_out}\")\n    except Exception as e:\n        print(f\"Error during Ensemble: {e}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- MAIN PIPELINE ---\nif __name__ == \"__main__\":\n    print(\"--- CAFA 6 PIPELINE START ---\")\n    \n    # Pre-Flight Check\n    if not os.path.exists(TRAIN_FASTA):\n        print(f\"CRITICAL ERROR: Input file not found: {TRAIN_FASTA}\")\n        print(\"Please ensure the 'cafa-6-protein-function-prediction' dataset is attached.\")\n    else:\n        try:\n            # 1. Embeddings\n            train_ids, X_train = generate_embeddings(TRAIN_FASTA, f\"{EMBED_DIR}/train\")\n            test_ids, X_test = generate_embeddings(TEST_FASTA, f\"{EMBED_DIR}/test\")\n            \n            # 2. Naive\n            f_naive = f\"{OUTPUT_DIR}/submission_naive.tsv\"\n            run_naive(test_ids, f_naive)\n            \n            # 3. ResNet\n            f_resnet = f\"{OUTPUT_DIR}/submission_resnet.tsv\"\n            train_predict_resnet(train_ids, X_train, test_ids, X_test, f_resnet)\n            \n            # 4. Similarity\n            f_sim = f\"{OUTPUT_DIR}/submission_similarity.tsv\"\n            has_sim = run_similarity(f_sim)\n            \n            # 5. Ensemble\n            f_final = f\"{OUTPUT_DIR}/submission.tsv\"\n            run_ensemble(f_naive, f_resnet, f_sim if has_sim else \"dummy\", f_final)\n            \n            print(\"--- PIPELINE COMPLETE: SUCCESS ---\")\n            \n        except Exception as e:\n            print(f\"--- PIPELINE FAILED ---\")\n            print(f\"Error: {e}\")\n            import traceback\n            traceback.print_exc()","metadata":{},"outputs":[],"execution_count":null}]}