{"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":14084779,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install biopython obonet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:28.887791Z","iopub.execute_input":"2025-10-19T18:18:28.888333Z","iopub.status.idle":"2025-10-19T18:18:35.428675Z","shell.execute_reply.started":"2025-10-19T18:18:28.888311Z","shell.execute_reply":"2025-10-19T18:18:35.427931Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nfrom Bio import SeqIO   # for FASTA parsing\nimport obonet           # for GO graph\n\n# paths (Kaggle input directory)\nDATA_PATH = \"/kaggle/input/cafa-6-protein-function-prediction\"\n\n# TSV files\nterms_df = pd.read_csv(f\"{DATA_PATH}/Train/train_terms.tsv\", sep=\"\\t\", names=[\"protein_id\", \"go_id\", \"ontology\"])\ntaxonomy_df = pd.read_csv(f\"{DATA_PATH}/Train/train_taxonomy.tsv\", sep=\"\\t\", names=[\"protein_id\", \"taxon_id\"])\nia_df = pd.read_csv(f\"{DATA_PATH}/IA.tsv\", sep=\"\\t\", names=[\"go_id\", \"ia_weight\"])\n\n# FASTA sequences\nsequences = []\nfor record in SeqIO.parse(f\"{DATA_PATH}/Train/train_sequences.fasta\", \"fasta\"):\n    sequences.append({\n        \"protein_id\": record.id.split(\"|\")[1],\n        \"sequence\": str(record.seq)\n    })\nseq_df = pd.DataFrame(sequences)\n\n# Merge together\ntrain_df = seq_df.merge(terms_df, on=\"protein_id\").merge(taxonomy_df, on=\"protein_id\", how=\"left\")\n\nprint(train_df.head())","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:35.430020Z","iopub.execute_input":"2025-10-19T18:18:35.430241Z","iopub.status.idle":"2025-10-19T18:18:37.860461Z","shell.execute_reply.started":"2025-10-19T18:18:35.430222Z","shell.execute_reply":"2025-10-19T18:18:37.859478Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Drop duplicates\ntrain_df = train_df.drop_duplicates()\n\n# Remove very short sequences (optional)\ntrain_df = train_df[train_df[\"sequence\"].str.len() > 30]\n\n# Check ontology distribution\nprint(train_df[\"ontology\"].value_counts())\n\nprotein2go = train_df.groupby(\"protein_id\")[\"go_id\"].apply(list).to_dict()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:37.861535Z","iopub.execute_input":"2025-10-19T18:18:37.861862Z","iopub.status.idle":"2025-10-19T18:18:39.976304Z","shell.execute_reply.started":"2025-10-19T18:18:37.861839Z","shell.execute_reply":"2025-10-19T18:18:39.975734Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n\nprotein_ids = list(protein2go.keys())\n\ntrain_ids, temp_ids = train_test_split(protein_ids, test_size=0.3, random_state=42)\nval_ids, test_ids = train_test_split(temp_ids, test_size=0.5, random_state=42)\n\nprint(f\"Train: {len(train_ids)} | Val: {len(val_ids)} | Test: {len(test_ids)}\")\ndel temp_ids","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:39.977827Z","iopub.execute_input":"2025-10-19T18:18:39.978082Z","iopub.status.idle":"2025-10-19T18:18:40.532222Z","shell.execute_reply.started":"2025-10-19T18:18:39.978064Z","shell.execute_reply":"2025-10-19T18:18:40.531636Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Sequence length distribution\ntrain_df[\"seq_len\"] = train_df[\"sequence\"].str.len()\nprint(train_df[\"seq_len\"].describe())\n\n# Remove extremely long or short sequences if needed (optional)\ntrain_df = train_df[(train_df[\"seq_len\"] >= 30) & (train_df[\"seq_len\"] <= 2000)]\n\n# Check average GO terms per protein\nterms_per_protein = train_df.groupby(\"protein_id\")[\"go_id\"].nunique()\nprint(terms_per_protein.describe())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:40.532862Z","iopub.execute_input":"2025-10-19T18:18:40.533192Z","iopub.status.idle":"2025-10-19T18:18:40.910756Z","shell.execute_reply.started":"2025-10-19T18:18:40.533176Z","shell.execute_reply":"2025-10-19T18:18:40.910073Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# protein → list of GO terms\nprotein2go = train_df.groupby(\"protein_id\")[\"go_id\"].apply(list).to_dict()\n\n# GO term → index\nall_terms = sorted(train_df[\"go_id\"].unique())\nterm2idx = {t: i for i, t in enumerate(all_terms)}\nidx2term = {i: t for t, i in term2idx.items()}\n\nprint(f\"Total unique GO terms: {len(all_terms)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:40.911535Z","iopub.execute_input":"2025-10-19T18:18:40.911777Z","iopub.status.idle":"2025-10-19T18:18:42.367449Z","shell.execute_reply.started":"2025-10-19T18:18:40.911759Z","shell.execute_reply":"2025-10-19T18:18:42.366472Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df_split = train_df[train_df[\"protein_id\"].isin(train_ids)]\nval_df_split   = train_df[train_df[\"protein_id\"].isin(val_ids)]\ntest_df_split  = train_df[train_df[\"protein_id\"].isin(test_ids)]\n\nprint(f\"Train rows: {len(train_df_split)}, Val rows: {len(val_df_split)}, Test rows: {len(test_df_split)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:42.368455Z","iopub.execute_input":"2025-10-19T18:18:42.368922Z","iopub.status.idle":"2025-10-19T18:18:42.560648Z","shell.execute_reply.started":"2025-10-19T18:18:42.368867Z","shell.execute_reply":"2025-10-19T18:18:42.559502Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# ontology graph","metadata":{}},{"cell_type":"code","source":"import networkx as nx\nimport obonet\n\n# Load ontology graph\ngo_graph = obonet.read_obo(f\"{DATA_PATH}/Train/go-basic.obo\")\n\n# Helper: get all ancestors for a term (excluding itself)\ndef get_ancestors(term_id):\n    if term_id in go_graph:\n        return nx.ancestors(go_graph, term_id)\n    return set()\n\n# Example\nexample_term = \"GO:0005634\"  # nucleus\nprint(\"Ancestors:\", get_ancestors(example_term))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:42.561982Z","iopub.execute_input":"2025-10-19T18:18:42.562293Z","iopub.status.idle":"2025-10-19T18:18:47.515264Z","shell.execute_reply.started":"2025-10-19T18:18:42.562265Z","shell.execute_reply":"2025-10-19T18:18:47.514469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# propagate GO terms for each protein\nprotein2go_propagated = {}\n\nfor pid, terms in protein2go.items():\n    expanded = set(terms)\n    for t in terms:\n        expanded.update(get_ancestors(t))\n    protein2go_propagated[pid] = list(expanded)\n\n# Convert to DataFrame (optional)\nexpanded_rows = []\nfor pid, terms in protein2go_propagated.items():\n    for t in terms:\n        expanded_rows.append((pid, t))\nexpanded_df = pd.DataFrame(expanded_rows, columns=[\"protein_id\", \"go_id\"])\nprint(expanded_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:18:47.515962Z","iopub.execute_input":"2025-10-19T18:18:47.516153Z","iopub.status.idle":"2025-10-19T18:20:14.492766Z","shell.execute_reply.started":"2025-10-19T18:18:47.516137Z","shell.execute_reply":"2025-10-19T18:20:14.492073Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧩 Step 9: Multi-Label Encoding (One-Hot or Multi-Hot Vectors)","metadata":{}},{"cell_type":"code","source":"import numpy as np\n\n# get all GO terms (vocabulary)\nall_terms = sorted(expanded_df[\"go_id\"].unique())\nterm2idx = {t: i for i, t in enumerate(all_terms)}\nidx2term = {i: t for t, i in term2idx.items()}\n\n# make vector for each protein\ndef encode_labels(term_list):\n    vec = np.zeros(len(all_terms), dtype=np.float32)\n    for t in term_list:\n        if t in term2idx:\n            vec[term2idx[t]] = 1.0\n    return vec\n\n# encode all\nprotein_vectors = {}\nfor pid, term_list in protein2go_propagated.items():\n    protein_vectors[pid] = encode_labels(term_list)\n\n# check one example\nexample_id = list(protein_vectors.keys())[0]\nprint(example_id)\nprint(np.sum(protein_vectors[example_id]), \"GO terms encoded as 1\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:20:14.494423Z","iopub.execute_input":"2025-10-19T18:20:14.494678Z","iopub.status.idle":"2025-10-19T18:20:44.097375Z","shell.execute_reply.started":"2025-10-19T18:20:14.494661Z","shell.execute_reply":"2025-10-19T18:20:44.096611Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧠 Step 10: Generate Protein Embeddings using ESM-2 (Facebook FAIR)","metadata":{}},{"cell_type":"code","source":"!pip install torch torchvision torchaudio\n!pip install fair-esm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:20:44.098254Z","iopub.execute_input":"2025-10-19T18:20:44.098641Z","iopub.status.idle":"2025-10-19T18:22:05.716204Z","shell.execute_reply.started":"2025-10-19T18:20:44.098614Z","shell.execute_reply":"2025-10-19T18:22:05.715469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Setup (run once) ---\nimport os, gc, pickle, numpy as np, torch, time\nfrom esm import pretrained\nfrom Bio import SeqIO\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\ntorch.set_float32_matmul_precision(\"high\")\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\nDATA_PATH = \"/kaggle/input/cafa-6-protein-function-prediction\"\n\n# ✅ Ensure the correct FASTA path (remove /Train if it doesn't exist)\nfasta_path = f\"{DATA_PATH}/Train/train_sequences.fasta\"\n\n# Load sequences\nseq_dict = {}\nfor rec in SeqIO.parse(fasta_path, \"fasta\"):\n    pid = rec.id.split(\"|\")[1]\n    seq_dict[pid] = str(rec.seq)\n\n# ✅ ESM2-T12-35M (last layer = 12, embedding dim = 480)\nesm_model, alphabet = pretrained.esm2_t12_35M_UR50D()\nbatch_converter = alphabet.get_batch_converter()\nesm_model = esm_model.to(device).eval()\n\nEMB_PATH = \"/kaggle/working/esm2_t12_train_embeds.pkl\"\n\n# Resume if exists\ntry:\n    with open(EMB_PATH, \"rb\") as f:\n        all_embeddings = pickle.load(f)\n    print(f\"Loaded existing embeddings: {len(all_embeddings)}\")\nexcept Exception:\n    all_embeddings = {}\n    print(\"No existing embedding file, starting fresh.\")\n\ndef chunked(lst, n):\n    for i in range(0, len(lst), n):\n        yield lst[i:i+n]\n\ndef generate_esm_embeddings_for_ids(\n    protein_ids,\n    max_len=1022,\n    batch_size=4,       # T12 is lighter; try 4, fall back to 2 or 1 if OOM\n    layer=12,           # ✅ T12 last layer\n    save_every=2000\n):\n    global all_embeddings\n    todo = [pid for pid in protein_ids if pid not in all_embeddings and pid in seq_dict]\n    # Sort by (truncated) length to reduce peaks\n    todo.sort(key=lambda x: min(len(seq_dict[x]), max_len))\n\n    start = time.time()\n    done_counter = 0\n    total = len(todo)\n\n    for group in chunked(todo, batch_size):\n        batch_pairs = []\n        for pid in group:\n            seq = seq_dict[pid]\n            if len(seq) > max_len:\n                seq = seq[:max_len]\n            batch_pairs.append((pid, seq))\n\n        labels, strs, toks = batch_converter(batch_pairs)\n        toks = toks.to(device)\n\n        with torch.no_grad():\n            with torch.amp.autocast(\"cuda\", dtype=torch.float16):  # ✅ new autocast API\n                out = esm_model(toks, repr_layers=[layer], return_contacts=False)\n        reps = out[\"representations\"][layer]\n\n        for j, (pid, seq) in enumerate(batch_pairs):\n            emb = reps[j, 1:1+len(seq)].mean(0).float().cpu().numpy()  # (480,)\n            all_embeddings[pid] = emb\n            done_counter += 1\n\n        # Cleanup\n        del labels, strs, toks, out, reps\n        torch.cuda.empty_cache(); gc.collect()\n\n        if done_counter % save_every == 0:\n            with open(EMB_PATH, \"wb\") as f:\n                pickle.dump(all_embeddings, f)\n            elapsed = time.time() - start\n            print(f\"Saved progress: {len(all_embeddings)} embeddings | \"\n                  f\"[{done_counter}/{total}] ~{done_counter/max(elapsed,1):.2f} prot/s\")\n\n    with open(EMB_PATH, \"wb\") as f:\n        pickle.dump(all_embeddings, f)\n    print(f\"Embedding complete. Total proteins embedded: {len(all_embeddings)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:22:05.717267Z","iopub.execute_input":"2025-10-19T18:22:05.717573Z","iopub.status.idle":"2025-10-19T18:22:14.277287Z","shell.execute_reply.started":"2025-10-19T18:22:05.717550Z","shell.execute_reply":"2025-10-19T18:22:14.276597Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ids_to_embed = list(set(train_ids) | set(val_ids))  # as you already had\ngenerate_esm_embeddings_for_ids(ids_to_embed, max_len=1022, batch_size=4, layer=12)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T09:40:18.528724Z","iopub.execute_input":"2025-10-19T09:40:18.529308Z","execution_failed":"2025-10-19T17:46:10.290Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ids_to_embed = list(set(train_ids) | set(val_ids))  # as you already had\ngenerate_esm_embeddings_for_ids(ids_to_embed, max_len=1022, batch_size=4, layer=12)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T18:43:34.868623Z","iopub.execute_input":"2025-10-19T18:43:34.869294Z","iopub.status.idle":"2025-10-19T20:08:24.844520Z","shell.execute_reply.started":"2025-10-19T18:43:34.869270Z","shell.execute_reply":"2025-10-19T20:08:24.843664Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🧠 Next Step — Step 11: Classifier Training (Multi-Label Model)","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\n\nclass ProteinGOClassifier(nn.Module):\n    def __init__(self, input_dim=1280, hidden_dim=512, output_dim=len(all_terms)):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(0.4),\n            nn.Linear(hidden_dim, output_dim)\n        )\n\n    def forward(self, x):\n        return self.net(x)  # raw logits\n\n# instantiate\nclf_model = ProteinGOClassifier().to(device)\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = optim.AdamW(clf_model.parameters(), lr=1e-4, weight_decay=1e-5)\n","metadata":{"trusted":true,"execution":{"execution_failed":"2025-10-19T22:58:23.460Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Dataset, DataLoader\nimport numpy as np\n\n# generate embeddings for all proteins (can take some time)\nall_embeddings = {}\nbatch_size = 4  # adjust depending on GPU memory\nprotein_list = list(seq_df[\"protein_id\"].values)\n\nfor i in range(0, len(protein_list), batch_size):\n    batch = protein_list[i:i+batch_size]\n    batch_seqs = [(pid, seq_df.loc[seq_df[\"protein_id\"] == pid, \"sequence\"].values[0]) for pid in batch]\n    batch_labels, batch_strs, batch_tokens = batch_converter(batch_seqs)\n    batch_tokens = batch_tokens.to(device)\n    with torch.no_grad():\n        out = esm_model(batch_tokens, repr_layers=[33], return_contacts=False)\n    reps = out[\"representations\"][33]\n    for j, (pid, seq) in enumerate(batch_seqs):\n        all_embeddings[pid] = reps[j, 1:len(seq)+1].mean(0).cpu().numpy()\n\nprint(\"Embeddings generated:\", len(all_embeddings))\n\nclass ProteinDataset(Dataset):\n    def __init__(self, pid_list, embeddings_dict, label_dict):\n        self.pids = pid_list\n        self.emb = embeddings_dict\n        self.lbl = label_dict\n    def __len__(self): return len(self.pids)\n    def __getitem__(self, idx):\n        pid = self.pids[idx]\n        return torch.tensor(self.emb[pid], dtype=torch.float32), \\\n               torch.tensor(self.lbl[pid], dtype=torch.float32)\n\ntrain_ds = ProteinDataset(train_ids, embeddings, protein_vectors)\nval_ds   = ProteinDataset(val_ids, embeddings, protein_vectors)\n\ntrain_loader = DataLoader(train_ds, batch_size=32, shuffle=True)\nval_loader   = DataLoader(val_ds, batch_size=32)\n\nfor epoch in range(5):\n    clf_model.train()\n    total_loss = 0\n    for x, y in train_loader:\n        x, y = x.to(device), y.to(device)\n        optimizer.zero_grad()\n        out = clf_model(x)\n        loss = criterion(out, y)\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n    print(f\"Epoch {epoch+1}, Train loss = {total_loss/len(train_loader):.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T05:55:53.615393Z","iopub.execute_input":"2025-10-19T05:55:53.616061Z","iopub.status.idle":"2025-10-19T05:57:22.533187Z","shell.execute_reply.started":"2025-10-19T05:55:53.616036Z","shell.execute_reply":"2025-10-19T05:57:22.531932Z"}},"outputs":[],"execution_count":null}]}