{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Packages","metadata":{}},{"cell_type":"code","source":"import warnings\nwarnings.filterwarnings('ignore')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-17T10:49:19.746976Z","iopub.execute_input":"2025-10-17T10:49:19.747219Z","iopub.status.idle":"2025-10-17T10:49:19.753752Z","shell.execute_reply.started":"2025-10-17T10:49:19.74718Z","shell.execute_reply":"2025-10-17T10:49:19.752898Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install biopython obonet --quiet\n!pip install torch torchvision torchaudio --quiet\n!pip install transformers biopython --quiet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-17T10:49:19.75457Z","iopub.execute_input":"2025-10-17T10:49:19.754834Z","iopub.status.idle":"2025-10-17T10:50:41.655144Z","shell.execute_reply.started":"2025-10-17T10:49:19.754812Z","shell.execute_reply":"2025-10-17T10:50:41.654256Z"},"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧩 CAFA6 Data Loading & Parsing (Step 1)","metadata":{}},{"cell_type":"code","source":"%%time\n# =========================================================\n# CAFA6 - Step 1: Load & Explore Data\n# =========================================================\nimport pandas as pd\n\nfrom Bio import SeqIO\nimport obonet\n\n# ---------------------------------------------------------\n# File paths (update if needed)\n# ---------------------------------------------------------\nBASE_PATH = \"/kaggle/input/cafa-6-protein-function-prediction\"\nTRAIN_PATH = f\"{BASE_PATH}/Train\"\nTEST_PATH = f\"{BASE_PATH}/Test\"\n\n# ---------------------------------------------------------\n# 1. Load TSV files\n# ---------------------------------------------------------\ntrain_terms = pd.read_csv(\n    f\"{TRAIN_PATH}/train_terms.tsv\",\n    sep=\"\\t\",\n    names=[\"protein_id\", \"go_term\", \"ontology\"]\n)\n\ntrain_taxonomy = pd.read_csv(\n    f\"{TRAIN_PATH}/train_taxonomy.tsv\",\n    sep=\"\\t\",\n    names=[\"protein_id\", \"taxon_id\"]\n)\n\nia = pd.read_csv(\n    f\"{BASE_PATH}/IA.tsv\",\n    sep=\"\\t\",\n    names=[\"go_term\", \"weight\"]\n)\n\n# ---------------------------------------------------------\n# 2. Load ontology graph (go-basic.obo)\n# ---------------------------------------------------------\ngo_graph = obonet.read_obo(f\"{TRAIN_PATH}/go-basic.obo\")\nprint(f\"Ontology terms loaded: {len(go_graph)}\")\n\n# Example: extract parent relationships\nroot_nodes = {\n    'BP': 'GO:0008150',\n    'CC': 'GO:0005575',\n    'MF': 'GO:0003674'\n}\n\n# ---------------------------------------------------------\n# 3. Load FASTA sequences (train & test superset)\n# ---------------------------------------------------------\ndef load_fasta_sequences(file_path):\n    \"\"\"Return dict {protein_id: sequence}\"\"\"\n    sequences = {}\n    for record in SeqIO.parse(file_path, \"fasta\"):\n        prot_id = record.id.split(\"|\")[1] if \"|\" in record.id else record.id\n        sequences[prot_id] = str(record.seq)\n    return sequences\n\ntrain_sequences = load_fasta_sequences(f\"{TRAIN_PATH}/train_sequences.fasta\")\ntest_sequences = load_fasta_sequences(f\"{TEST_PATH}/testsuperset.fasta\")\n\nprint(f\"Train sequences: {len(train_sequences)}\")\nprint(f\"Test superset sequences: {len(test_sequences)}\")\n\n# ---------------------------------------------------------\n# 4. Merge annotations and taxonomy\n# ---------------------------------------------------------\ntrain_df = train_terms.merge(train_taxonomy, on=\"protein_id\", how=\"left\")\nprint(f\"Train dataframe: {train_df.shape}\")\n\n# ---------------------------------------------------------\n# 5. Quick data summary\n# ---------------------------------------------------------\nprint(\"\\nData Summary:\")\nprint(f\"- train_terms: {train_terms.shape}\")\nprint(f\"- train_taxonomy: {train_taxonomy.shape}\")\nprint(f\"- ia: {ia.shape}\")\nprint(f\"- train_sequences: {len(train_sequences)}\")\nprint(f\"- test_sequences: {len(test_sequences)}\")\nprint(f\"- Ontology terms (GO): {len(go_graph)}\")\n\n# Preview few examples\nprint(\"\\nSample train_terms:\")\ndisplay(train_terms.head())\n\nprint(\"\\nSample train_taxonomy:\")\ndisplay(train_taxonomy.head())\n\nprint(\"\\nSample IA:\")\ndisplay(ia.head())\n\n# Example check of GO hierarchy\nexample_term = list(go_graph.nodes())[0]\nprint(f\"\\nExample GO term: {example_term}\")\nprint(\"Parents:\", list(go_graph.predecessors(example_term)))\nprint(\"Children:\", list(go_graph.successors(example_term)))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-17T10:50:41.657497Z","iopub.execute_input":"2025-10-17T10:50:41.657798Z","iopub.status.idle":"2025-10-17T10:50:50.110645Z","shell.execute_reply.started":"2025-10-17T10:50:41.657775Z","shell.execute_reply":"2025-10-17T10:50:50.110054Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧭 Step 2: Exploratory Data Analysis (EDA)\n\nBelow is a well-structured EDA section to add right after dataset loading.\n\nIt covers:\n\nOverview of GO ontologies\n\nDistribution of GO terms & taxonomies\n\nSequence length statistics\n\nRelationship between ontology, taxonomy, and terms\n\nVisualization of embeddings (optional)","metadata":{}},{"cell_type":"code","source":"%%time\n# =========================================================\n# Step 2: Exploratory Data Analysis (EDA)\n# =========================================================\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\n\n# Style setup\nplt.style.use(\"seaborn-v0_8-muted\")\nsns.set_palette(\"coolwarm\")\n\n# ---------------------------------------------------------\n# 1. Ontology distribution\n# ---------------------------------------------------------\nplt.figure(figsize=(6,4))\nsns.countplot(x=\"ontology\", data=train_terms)\nplt.title(\"Distribution of Ontology Types\")\nplt.xlabel(\"Ontology (Aspect)\")\nplt.ylabel(\"Count\")\nplt.tight_layout()\nplt.show()\n\n# ---------------------------------------------------------\n# 2. GO term frequency\n# ---------------------------------------------------------\ngo_counts = train_terms[\"go_term\"].value_counts().head(20)\nplt.figure(figsize=(8,4))\nsns.barplot(y=go_counts.index, x=go_counts.values)\nplt.title(\"Top 20 Most Frequent GO Terms\")\nplt.xlabel(\"Count\")\nplt.ylabel(\"GO Term\")\nplt.tight_layout()\nplt.show()\n\n# ---------------------------------------------------------\n# 3. Taxonomy overview\n# ---------------------------------------------------------\ntax_counts = train_taxonomy[\"taxon_id\"].value_counts().head(10)\nplt.figure(figsize=(8,4))\nsns.barplot(x=tax_counts.index.astype(str), y=tax_counts.values)\nplt.title(\"Top 10 Taxon IDs (Species)\")\nplt.xlabel(\"Taxon ID\")\nplt.ylabel(\"Protein Count\")\nplt.tight_layout()\nplt.show()\n\n# ---------------------------------------------------------\n# 4. Sequence length distribution\n# ---------------------------------------------------------\nseq_lengths = [len(seq) for seq in train_sequences.values()]\nplt.figure(figsize=(8,4))\nsns.histplot(seq_lengths, bins=50, kde=True)\nplt.title(\"Distribution of Protein Sequence Lengths\")\nplt.xlabel(\"Sequence Length (aa)\")\nplt.ylabel(\"Count\")\nplt.tight_layout()\nplt.show()\n\nprint(f\"🧬 Average sequence length: {np.mean(seq_lengths):.1f} aa\")\nprint(f\"Longest sequence: {np.max(seq_lengths)} aa\")\nprint(f\"Shortest sequence: {np.min(seq_lengths)} aa\")\n\n# ---------------------------------------------------------\n# 5. Ontology × Taxonomy heatmap (co-occurrence)\n# ---------------------------------------------------------\nmerged = train_terms.merge(train_taxonomy, on=\"protein_id\", how=\"left\")\nont_tax_counts = (\n    merged.groupby([\"ontology\", \"taxon_id\"])\n    .size()\n    .reset_index(name=\"count\")\n    .pivot(index=\"ontology\", columns=\"taxon_id\", values=\"count\")\n    .fillna(0)\n)\nplt.figure(figsize=(10,4))\nsns.heatmap(np.log1p(ont_tax_counts), cmap=\"viridis\")\nplt.title(\"Ontology × Taxonomy Co-occurrence (log scale)\")\nplt.xlabel(\"Taxon ID\")\nplt.ylabel(\"Ontology\")\nplt.tight_layout()\nplt.show()\n\n# ---------------------------------------------------------\n# 6. IA weight distribution (Information Content)\n# ---------------------------------------------------------\nplt.figure(figsize=(8,4))\nsns.histplot(ia[\"weight\"], bins=50, kde=True)\nplt.title(\"GO Term Information Content (IA Weights)\")\nplt.xlabel(\"Weight\")\nplt.ylabel(\"Frequency\")\nplt.tight_layout()\nplt.show()\n\nprint(f\"🧠 IA weights: mean={ia['weight'].mean():.3f}, max={ia['weight'].max():.3f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-17T10:50:50.111293Z","iopub.execute_input":"2025-10-17T10:50:50.111468Z","iopub.status.idle":"2025-10-17T10:50:53.614541Z","shell.execute_reply.started":"2025-10-17T10:50:50.111452Z","shell.execute_reply":"2025-10-17T10:50:53.613919Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧬 Step 3 – Protein Embedding Extraction (ESM2 baseline)","metadata":{}},{"cell_type":"code","source":"%%time\n# =========================================================\n# Step 2: Extract Protein Embeddings using ESM2 (facebook/esm2_t6_8M_UR50D)\n# =========================================================\n\n\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\nfrom Bio import SeqIO\nimport numpy as np\nimport pandas as pd\nfrom tqdm.notebook import tqdm\n\n# ---------------------------------------------------------\n# Load model + tokenizer\n# ---------------------------------------------------------\nMODEL_NAME = \"facebook/esm2_t6_8M_UR50D\"  # small and fast baseline\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModel.from_pretrained(MODEL_NAME)\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel = model.to(device)\nmodel.eval()\n\nprint(f\"✅ Loaded {MODEL_NAME} on {device}\")\n\n# ---------------------------------------------------------\n# Function to compute mean-pooled embeddings\n# ---------------------------------------------------------\ndef get_protein_embedding(sequence: str):\n    \"\"\"Return mean pooled embedding for one protein sequence.\"\"\"\n    inputs = tokenizer(sequence, return_tensors=\"pt\", truncation=True, max_length=1022)\n    with torch.no_grad():\n        outputs = model(**{k: v.to(device) for k, v in inputs.items()})\n    # Mean-pool across tokens (excluding [CLS], [EOS])\n    emb = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()\n    return emb\n\n# ---------------------------------------------------------\n# Compute embeddings for all training sequences\n# ---------------------------------------------------------\n# You can limit to first N for quick testing\nN = 1000  # change to None or len(train_sequences) for full run\n\nseq_items = list(train_sequences.items())[:N]\nembeddings = []\nprotein_ids = []\n\nfor prot_id, seq in tqdm(seq_items, desc=\"Embedding proteins\"):\n    try:\n        emb = get_protein_embedding(seq)\n        embeddings.append(emb)\n        protein_ids.append(prot_id)\n    except Exception as e:\n        print(f\"⚠️ Skipped {prot_id}: {e}\")\n\n# ---------------------------------------------------------\n# Save embeddings\n# ---------------------------------------------------------\nemb_df = pd.DataFrame(embeddings)\nemb_df.insert(0, \"protein_id\", protein_ids)\nemb_df.to_parquet(\"train_esm2_embeddings.parquet\", index=False)\n\nprint(f\"\\n✅ Saved {len(emb_df)} protein embeddings with shape {emb_df.shape}\")\ndisplay(emb_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-17T10:50:53.615361Z","iopub.execute_input":"2025-10-17T10:50:53.615863Z","iopub.status.idle":"2025-10-17T10:51:40.154255Z","shell.execute_reply.started":"2025-10-17T10:50:53.615832Z","shell.execute_reply":"2025-10-17T10:51:40.153473Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# ⚙️ Step 4 – Multi-Label Classifier for GO Term Prediction\n\nThis step links:\n\nthe embeddings (emb_df)\n\nthe GO labels (train_terms)\n\nThen trains a simple baseline model.","metadata":{}},{"cell_type":"code","source":"%%time\n# =========================================================\n# Step 3: Multi-label GO-term classification (Baseline)\n# =========================================================\n!pip install scikit-learn --quiet\nimport pandas as pd\nimport numpy as np\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score\nfrom tqdm.notebook import tqdm\n\n# ---------------------------------------------------------\n# 1. Load embeddings and labels\n# ---------------------------------------------------------\nemb_df = pd.read_parquet(\"train_esm2_embeddings.parquet\")\n\n# Filter train_terms to proteins we have embeddings for\nsubset_terms = train_terms[train_terms[\"protein_id\"].isin(emb_df[\"protein_id\"])]\n\n# Merge ontology info if needed\nsubset_terms = subset_terms.merge(train_taxonomy, on=\"protein_id\", how=\"left\")\n\nprint(f\"Training proteins used: {subset_terms['protein_id'].nunique()}\")\n\n# ---------------------------------------------------------\n# 2. Prepare multi-label targets\n# ---------------------------------------------------------\n# Keep top K frequent GO terms for faster training\nK = 50\ntop_terms = subset_terms[\"go_term\"].value_counts().head(K).index\nsubset_terms = subset_terms[subset_terms[\"go_term\"].isin(top_terms)]\n\n# Create label list per protein\nlabels_df = subset_terms.groupby(\"protein_id\")[\"go_term\"].apply(list).reset_index()\n\n# Join embeddings\ntrain_data = emb_df.merge(labels_df, on=\"protein_id\", how=\"inner\")\nX = train_data.drop(columns=[\"protein_id\", \"go_term\"]).values\ny_list = train_data[\"go_term\"].tolist()\n\nmlb = MultiLabelBinarizer()\nY = mlb.fit_transform(y_list)\n\nprint(f\"Feature matrix: {X.shape}, Label matrix: {Y.shape}\")\n\n# ---------------------------------------------------------\n# 3. Train/Test split\n# ---------------------------------------------------------\nX_train, X_val, y_train, y_val = train_test_split(X, Y, test_size=0.2, random_state=42)\n\n# ---------------------------------------------------------\n# 4. Train One-vs-Rest Logistic Regression\n# ---------------------------------------------------------\nclf = OneVsRestClassifier(LogisticRegression(max_iter=200))\nclf.fit(X_train, y_train)\n\n# ---------------------------------------------------------\n# 5. Evaluate\n# ---------------------------------------------------------\ny_pred = clf.predict(X_val)\nf1 = f1_score(y_val, y_pred, average=\"micro\")\n\nprint(f\"\\n✅ Baseline micro-F1 score on validation set: {f1:.4f}\")\n\n# Inspect example predictions\nexample_idx = np.random.randint(0, len(y_pred))\npred_terms = [mlb.classes_[i] for i, v in enumerate(y_pred[example_idx]) if v == 1]\nprint(f\"\\nExample protein predicted GO terms:\\n{pred_terms}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-17T10:51:40.155103Z","iopub.execute_input":"2025-10-17T10:51:40.155677Z","iopub.status.idle":"2025-10-17T10:51:59.724702Z","shell.execute_reply.started":"2025-10-17T10:51:40.155656Z","shell.execute_reply":"2025-10-17T10:51:59.723814Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧾 Step 5 – Submission File Generator (CAFA6 format)\n\nThis will:\n\npredict GO term probabilities for all proteins you choose (train/test),\n\nkeep top predictions per protein (up to 1500 total per rules),\n\nand write a valid tab-separated file.","metadata":{}},{"cell_type":"code","source":"%%time\n# =========================================================\n# Step 4: Submission File Generator\n# =========================================================\nimport numpy as np\nimport pandas as pd\n\n# ---------------------------------------------------------\n# 1. Predict probabilities on all available embeddings\n# ---------------------------------------------------------\nproba = clf.predict_proba(X)  # probability scores\npred_df = pd.DataFrame(proba, columns=mlb.classes_)\npred_df.insert(0, \"protein_id\", train_data[\"protein_id\"].values)\n\n# ---------------------------------------------------------\n# 2. Convert to long format (protein_id, GO_term, score)\n# ---------------------------------------------------------\nsubmissions = []\nfor idx, row in pred_df.iterrows():\n    pid = row[\"protein_id\"]\n    scores = row[1:].values\n    for go_term, score in zip(pred_df.columns[1:], scores):\n        if score > 0.01:  # keep only confident predictions\n            submissions.append((pid, go_term, round(float(score), 3)))\n\nsubmission_df = pd.DataFrame(submissions, columns=[\"protein_id\", \"go_term\", \"score\"])\n\n# ---------------------------------------------------------\n# 3. Enforce CAFA6 constraints\n# ---------------------------------------------------------\n# - Keep max 1500 terms per protein\nsubmission_df = (\n    submission_df.sort_values([\"protein_id\", \"score\"], ascending=[True, False])\n    .groupby(\"protein_id\")\n    .head(1500)\n)\n\n# - Ensure score in (0, 1.000]\nsubmission_df = submission_df[submission_df[\"score\"] > 0]\n\n# ---------------------------------------------------------\n# 4. Save to tsv (no header)\n# ---------------------------------------------------------\nsubmission_df.to_csv(\"submission.tsv\", sep=\"\\t\", index=False, header=False)\n\nprint(f\"✅ Submission file ready: {submission_df.shape[0]} lines\")\ndisplay(submission_df.head(10))","metadata":{"trusted":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-10-17T10:51:59.725953Z","iopub.execute_input":"2025-10-17T10:51:59.726216Z","iopub.status.idle":"2025-10-17T10:52:00.02247Z","shell.execute_reply.started":"2025-10-17T10:51:59.726172Z","shell.execute_reply":"2025-10-17T10:52:00.021638Z"}},"outputs":[],"execution_count":null}]}