{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14154175,"sourceType":"datasetVersion","datasetId":8932669}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CAFA-6 Protein Function Prediction: The Ensemble Era \n\nWelcome to this comprehensive pipeline for the **CAFA-6 Competition**!\n\nIn this notebook, we will build a powerful, **multi-model ensemble** system to predict protein functions (GO terms) from amino acid sequences. We will combine the strengths of three state-of-the-art Protein Language Models (pLMs):\n\n1.  **ProtBERT**  (Rostlab)\n2.  **ESM-2**  (Meta AI)\n\n###  Key Features\n-   **Visual EDA**: Beautiful plots to understand our data.\n-   **GPU Optimization**: Efficient embedding generation.\n-   **Ensemble Learning**: Averaging predictions for robustness.\n-   **Clean Code**: Structured, commented, and easy to follow.\n\nLet's dive in! ","metadata":{}},{"cell_type":"code","source":"!pip install biopython --quiet # I was getting an error so that's why doing this","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:01:41.555118Z","iopub.execute_input":"2025-12-15T13:01:41.555273Z","iopub.status.idle":"2025-12-15T13:01:47.002894Z","shell.execute_reply.started":"2025-12-15T13:01:41.555258Z","shell.execute_reply":"2025-12-15T13:01:47.001976Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Imports & Setup\nimport os\nimport sys\nimport gc\nimport json\nimport re\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom transformers import AutoTokenizer, AutoModel, T5EncoderModel\nfrom Bio import SeqIO\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom scipy import sparse\nfrom torch.utils.data import TensorDataset, DataLoader\nfrom multiprocessing import Pool, cpu_count\nimport glob\n# TPU SUPPORT (ADDED\nUSE_TPU = False\ntry:\n    import torch_xla\n    import torch_xla.core.xla_model as xm\n    import torch_xla.distributed.parallel_loader as pl\n    import torch_xla.debug.metrics as met\n    import torch_xla.utils.utils as xu\n    USE_TPU = True\nexcept ImportError:\n    USE_TPU = False\n\n#  Plotting Style\nsns.set_theme(style=\"whitegrid\", palette=\"viridis\")\nplt.rcParams['figure.figsize'] = (12, 6)\nplt.rcParams['font.size'] = 12\n\n#  Paths\nINPUT_DIR = \"/kaggle/input/cafa-6-protein-function-prediction/\"\nWORKING_DIR = \"/kaggle/working/\"\nOUTPUT_LABELS_DIR = os.path.join(WORKING_DIR, \"outputs_labels\")\nEMBEDDINGS_DIR = os.path.join(WORKING_DIR, \"embeddings\")\nMODELS_DIR = os.path.join(WORKING_DIR, \"models\")\nos.makedirs(OUTPUT_LABELS_DIR, exist_ok=True)\nos.makedirs(EMBEDDINGS_DIR, exist_ok=True)\nos.makedirs(MODELS_DIR, exist_ok=True)\n\nif USE_TPU:\n    device = torch_xla.device()\n    print(\" TPU detected. Using device:\", device)\nelse:\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(\" Using device:\", device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:01:47.005208Z","iopub.execute_input":"2025-12-15T13:01:47.005445Z","iopub.status.idle":"2025-12-15T13:02:15.427325Z","shell.execute_reply.started":"2025-12-15T13:01:47.005404Z","shell.execute_reply":"2025-12-15T13:02:15.426566Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Part 1: Exploratory Data Analysis (EDA)\nBefore we model, we must understand. Let's load the data and visualize the sequence lengths and GO term distributions.","metadata":{}},{"cell_type":"code","source":"#  Load Data\nprint(\" Loading sequences...\")\ntrain_seqs = {rec.id: str(rec.seq) for rec in SeqIO.parse(os.path.join(INPUT_DIR, \"Train/train_sequences.fasta\"), \"fasta\")}\ntest_seqs = {rec.id: str(rec.seq) for rec in SeqIO.parse(os.path.join(INPUT_DIR, \"Test/testsuperset.fasta\"), \"fasta\")}\ntrain_terms_df = pd.read_csv(os.path.join(INPUT_DIR, \"Train/train_terms.tsv\"), sep=\"\\t\")\n\nprint(f\" Train Sequences: {len(train_seqs):,}\")\nprint(f\" Test Sequences: {len(test_seqs):,}\")\nprint(f\" Train Annotations: {len(train_terms_df):,}\")\n\n# Plot Sequence Length Distribution\ntrain_lens = [len(s) for s in train_seqs.values()]\ntest_lens = [len(s) for s in test_seqs.values()]\n\nplt.figure(figsize=(14, 6))\nsns.histplot(train_lens, color=\"blue\", label=\"Train\", kde=True, alpha=0.5, log_scale=True)\nsns.histplot(test_lens, color=\"orange\", label=\"Test\", kde=True, alpha=0.5, log_scale=True)\nplt.title(\" Sequence Length Distribution (Log Scale)\")\nplt.xlabel(\"Sequence Length\")\nplt.legend()\nplt.show()\n\n#  Plot Top GO Terms\ntop_terms = train_terms_df['term'].value_counts().head(20)\nplt.figure(figsize=(14, 8))\nsns.barplot(x=top_terms.values, y=top_terms.index, palette=\"viridis\")\nplt.title(\" Top 20 Most Frequent GO Terms\")\nplt.xlabel(\"Count\")\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:15.428165Z","iopub.execute_input":"2025-12-15T13:02:15.428728Z","iopub.status.idle":"2025-12-15T13:02:22.019872Z","shell.execute_reply.started":"2025-12-15T13:02:15.428701Z","shell.execute_reply":"2025-12-15T13:02:22.019239Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  Part 2: The Knowledge Graph (Gene Ontology)\nProteins functions are related! If a protein does X, and X is a subclass of Y, it also does Y. We need to parse the **Gene Ontology (GO)** graph to enforce these rules.","metadata":{}},{"cell_type":"code","source":"#  Parse OBO & Build Ancestors\ndef parse_obo(obo_file):\n    terms = {}\n    current_term = None\n    with open(obo_file, 'r') as f:\n        for line in f:\n            line = line.strip()\n            if line == \"[Term]\":\n                if current_term:\n                    terms[current_term['id']] = current_term\n                current_term = {'is_a': [], 'namespace': ''}\n            elif line.startswith(\"id: \"):\n                current_term['id'] = line[4:]\n            elif line.startswith(\"namespace: \"):\n                current_term['namespace'] = line[11:]\n            elif line.startswith(\"is_a: \"):\n                current_term['is_a'].append(line[6:].split(' ! ')[0])\n    if current_term:\n        terms[current_term['id']] = current_term\n    return terms\n\ngo_terms = parse_obo(os.path.join(INPUT_DIR, \"Train/go-basic.obo\"))\nprint(f\" Parsed {len(go_terms):,} GO terms.\")\n\n# Build Ancestors (Transitive Closure)\nancestors = {}\nfor term_id in tqdm(go_terms, desc=\" Building Graph\"):\n    queue = [term_id]\n    visited = set()\n    while queue:\n        curr = queue.pop(0)\n        if curr in visited: continue\n        visited.add(curr)\n        if curr in go_terms:\n            queue.extend(go_terms[curr]['is_a'])\n    ancestors[term_id] = list(visited)\nwith open(os.path.join(OUTPUT_LABELS_DIR, \"ancestors.json\"), \"w\") as f:\n    json.dump(ancestors, f)\nprint(\" Ancestor graph saved.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:22.020532Z","iopub.execute_input":"2025-12-15T13:02:22.020760Z","iopub.status.idle":"2025-12-15T13:02:22.939911Z","shell.execute_reply.started":"2025-12-15T13:02:22.020735Z","shell.execute_reply":"2025-12-15T13:02:22.939274Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_uniprot_id(header):\n    # Extract between first pair of pipes → sp|XXXX| or tr|XXXX|\n    m = re.match(r\"^[a-z]{2}\\|([^|]+)\\|\", header)\n    if m:\n        return m.group(1)\n    return header  # fallback\n\n# Build mapping: extract UniProt ID for every FASTA protein\ntrain_proteins_raw = list(train_seqs.keys())\ntrain_proteins = [extract_uniprot_id(h) for h in train_proteins_raw]\n\nprot_map = {pid: i for i, pid in enumerate(train_proteins)}\n\n# Filter annotation proteins to the ones that appear in FASTA\nvalid_annots = train_terms_df[train_terms_df[\"EntryID\"].isin(prot_map)]\n\nprint(\"Proteins in FASTA:\", len(train_seqs))\nprint(\"Proteins in annotations:\", len(train_terms_df))\nprint(\"Overlap:\", len(valid_annots))\n\nused_terms = set(train_terms_df[\"term\"].unique())\n\nterms_MF = [t for t in used_terms if t in go_terms and go_terms[t]['namespace'] == 'molecular_function']\nterms_BP = [t for t in used_terms if t in go_terms and go_terms[t]['namespace'] == 'biological_process']\nterms_CC = [t for t in used_terms if t in go_terms and go_terms[t]['namespace'] == 'cellular_component']\n\n\nterm_maps = {\n    'MF': {str(i): t for i, t in enumerate(terms_MF)},\n    'BP': {str(i): t for i, t in enumerate(terms_BP)},\n    'CC': {str(i): t for i, t in enumerate(terms_CC)}\n}\n\nwith open(os.path.join(OUTPUT_LABELS_DIR, \"maps.json\"), \"w\") as f:\n    json.dump(term_maps, f)\n\nfor ns, terms_list in [('MF', terms_MF), ('BP', terms_BP), ('CC', terms_CC)]:\n\n    go_to_idx = {t: i for i, t in enumerate(terms_list)}\n\n    df_ns = valid_annots[valid_annots[\"term\"].isin(terms_list)]\n    \n    p_idx = df_ns[\"EntryID\"].map(prot_map).values\n    t_idx = df_ns[\"term\"].map(go_to_idx).values\n\n    mat = sparse.coo_matrix(\n        (np.ones(len(p_idx)), (p_idx, t_idx)),\n        shape=(len(train_proteins), len(terms_list)),\n        dtype=np.int8\n    ).tocsr()\n\n    sparse.save_npz(os.path.join(OUTPUT_LABELS_DIR, f\"labels_{ns}.npz\"), mat)\n\n    print(f\"{ns}: proteins={mat.shape[0]}, terms={mat.shape[1]}, nonzero={mat.nnz}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:22.940621Z","iopub.execute_input":"2025-12-15T13:02:22.940835Z","iopub.status.idle":"2025-12-15T13:02:23.812322Z","shell.execute_reply.started":"2025-12-15T13:02:22.940810Z","shell.execute_reply":"2025-12-15T13:02:23.811556Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#cheking the .npz files\nfor ns in [\"MF\", \"BP\", \"CC\"]:\n    mat = sparse.load_npz(os.path.join(OUTPUT_LABELS_DIR, f\"labels_{ns}.npz\"))\n    print(ns, \"nonzero:\", mat.count_nonzero(), \"shape:\", mat.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.813730Z","iopub.execute_input":"2025-12-15T13:02:23.813959Z","iopub.status.idle":"2025-12-15T13:02:23.844162Z","shell.execute_reply.started":"2025-12-15T13:02:23.813938Z","shell.execute_reply":"2025-12-15T13:02:23.843636Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  Part 3: Feature Engineering (Embeddings)\nHere I will now extract rich features from protein sequences using **three** different models. This is the secret sauce! \nWe define a generic embedding function to reuse for all models.","metadata":{}},{"cell_type":"code","source":"def embed_sequences(model_name, seq_dict, batch_size=16, max_len=1024):\n\n    # FORCE GPU ONLY\n    if not torch.cuda.is_available():\n        raise SystemExit(\" GPU NOT FOUND — Cannot embed on CPU. Aborting as requested.\")\n\n    device = torch.device(\"cuda\")\n    print(\" FORCING GPU:\", device)\n\n    name = model_name.lower()\n    tokenizer = AutoTokenizer.from_pretrained(model_name)\n    model = AutoModel.from_pretrained(\n            model_name,\n            torch_dtype=torch.float16\n        )\n\n    model = model.to(device)\n    model.eval()\n\n    seqs = list(seq_dict.values())\n    ids = list(seq_dict.keys())\n    embeddings = []\n\n    for i in tqdm(range(0, len(seqs), batch_size)):\n        batch = seqs[i:i + batch_size]\n\n        # TOKENIZATION RULES\n        if \"t5\" in name:\n            batch_tok = [\" \".join(list(s)) for s in batch]\n        elif \"bert\" in name:\n            batch_tok = [\" \".join(list(s)) for s in batch]\n        else:\n            batch_tok = batch   \n\n        inputs = tokenizer(\n            batch_tok,\n            return_tensors=\"pt\",\n            padding=True,\n            truncation=True,\n            max_length=max_len\n        ).to(device)\n\n        with torch.no_grad():\n            with torch.cuda.amp.autocast(dtype=torch.float16):\n\n                outputs = model(**inputs)\n                hidden = outputs.last_hidden_state   \n                mask = inputs[\"attention_mask\"].unsqueeze(-1).float()\n\n                # mean pooling\n                emb = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)\n\n        # move to CPU only at the last moment\n        embeddings.append(emb.float().cpu().numpy())\n\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    return np.concatenate(embeddings), ids\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.846351Z","iopub.execute_input":"2025-12-15T13:02:23.846621Z","iopub.status.idle":"2025-12-15T13:02:23.858775Z","shell.execute_reply.started":"2025-12-15T13:02:23.846603Z","shell.execute_reply":"2025-12-15T13:02:23.858065Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**Note: I have done the embedding locally so you can check the dataset**","metadata":{}},{"cell_type":"code","source":"# emb_train_bert, _ = embed_sequences(\"Rostlab/prot_bert\", train_seqs, batch_size=128)\n# np.save(\"train_protbert.npy\", emb_train_bert)\n# del emb_train_bert\n\n# emb_test_bert, test_ids = embed_sequences(\"Rostlab/prot_bert\", test_seqs, batch_size=128)\n# np.save(\"test_protbert.npy\", emb_test_bert)\n# np.save(\"test_ids.npy\", test_ids)\n# del emb_test_bert\nprint(\"Protbert embeddings saved\") ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.859448Z","iopub.execute_input":"2025-12-15T13:02:23.859647Z","iopub.status.idle":"2025-12-15T13:02:23.872558Z","shell.execute_reply.started":"2025-12-15T13:02:23.859631Z","shell.execute_reply":"2025-12-15T13:02:23.871872Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# emb_train_esm, _ = embed_sequences(\"facebook/esm2_t30_150M_UR50D\", train_seqs, batch_size=8)\n# np.save(\"train_esm2.npy\", emb_train_esm)\n# del emb_train_esm\n\n# emb_test_esm, _ = embed_sequences(\"facebook/esm2_t30_150M_UR50D\", test_seqs, batch_size=8)\n# np.save(\"test_esm2.npy\", emb_test_esm)\n# del emb_test_esm\nprint(\"ESM-2(150M) embeddings saved\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.873211Z","iopub.execute_input":"2025-12-15T13:02:23.873398Z","iopub.status.idle":"2025-12-15T13:02:23.891041Z","shell.execute_reply.started":"2025-12-15T13:02:23.873384Z","shell.execute_reply":"2025-12-15T13:02:23.890398Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**I am using a different embedding function for Prot5 cause I was finding difficulty to create a single unifies one for Prott5**","metadata":{}},{"cell_type":"markdown","source":"**We freeze ProtT5 and generate contextual protein embeddings, then train lightweight downstream heads.**","metadata":{}},{"cell_type":"code","source":"# model_name = \"Rostlab/prot_t5_xl_uniref50\"\n# tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, do_lower_case=False)\n\n# def preprocess_for_t5(seq):\n#     return \" \".join(list(seq))\n\n# def tokenize_all(seqs, max_len=1024):\n#     tok_inputs = []\n#     for s in tqdm(seqs, desc=\"Pre-tokenizing\"):\n#         t = tokenizer(\n#             preprocess_for_t5(s),\n#             truncation=True,\n#             max_length=max_len,\n#             padding=\"max_length\",   # <- CRITICAL for speed\n#             return_tensors=\"np\"\n#         )\n#         tok_inputs.append((t[\"input_ids\"][0], t[\"attention_mask\"][0]))\n#     return np.array([x[0] for x in tok_inputs]), np.array([x[1] for x in tok_inputs])\n\n# train_ids_np, train_mask_np = tokenize_all(list(train_seqs.values()))\n# np.save(\"train_ids.npy\", train_ids_np)\n# np.save(\"train_mask.npy\", train_mask_np)\nprint(\"Saved: train_ids.npy, train_mask.npy\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.891641Z","iopub.execute_input":"2025-12-15T13:02:23.891892Z","iopub.status.idle":"2025-12-15T13:02:23.907629Z","shell.execute_reply.started":"2025-12-15T13:02:23.891872Z","shell.execute_reply":"2025-12-15T13:02:23.906912Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# def tokenize_all(seqs, max_len=1024):\n#     tok_inputs = []\n#     for s in tqdm(seqs, desc=\"Pre-tokenizing\"):\n#         t = tokenizer(\n#             preprocess_for_t5(s),\n#             truncation=True,\n#             max_length=max_len,\n#             padding=\"max_length\",\n#             return_tensors=\"np\"\n#         )\n#         tok_inputs.append((t[\"input_ids\"][0], t[\"attention_mask\"][0]))\n#     return np.array([x[0] for x in tok_inputs]), np.array([x[1] for x in tok_inputs])\n\n# test_ids_np, test_mask_np = tokenize_all(list(test_seqs.values()))\n\n# np.save(\"test_ids.npy\", test_ids_np)\n# np.save(\"test_mask.npy\", test_mask_np)\n\nprint(\"Saved: test_ids.npy, test_mask.npy\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.908447Z","iopub.execute_input":"2025-12-15T13:02:23.908982Z","iopub.status.idle":"2025-12-15T13:02:23.922841Z","shell.execute_reply.started":"2025-12-15T13:02:23.908963Z","shell.execute_reply":"2025-12-15T13:02:23.922242Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_ids = np.load(\"/kaggle/input/embeddings-for-cafa/train_ids_T5.npy\")\ntrain_mask = np.load(\"/kaggle/input/embeddings-for-cafa/train_mask.npy\")\ntest_ids = np.load(\"/kaggle/input/embeddings-for-cafa/test_ids_T5.npy\")\ntest_mask = np.load(\"/kaggle/input/embeddings-for-cafa/test_mask.npy\")\nprint('loaded the npy files necessary to embed T5')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:23.923497Z","iopub.execute_input":"2025-12-15T13:02:23.923778Z","iopub.status.idle":"2025-12-15T13:02:43.548983Z","shell.execute_reply.started":"2025-12-15T13:02:23.923757Z","shell.execute_reply":"2025-12-15T13:02:43.548328Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def prot_t5_embed(\n    model_name: str,\n    ids_np: np.ndarray,\n    mask_np: np.ndarray,\n    batch_size: int = 32,\n    layer_idx: int = 12,\n):\n\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    # Load model\n    model = T5EncoderModel.from_pretrained(\n        model_name,\n        torch_dtype=torch.float16 if device.type == \"cuda\" else torch.float32\n    ).to(device)\n\n    model.eval()\n\n    N = ids_np.shape[0]\n    hidden_dim = model.config.d_model\n\n    # Preallocate output\n    embeddings = np.zeros((N, hidden_dim), dtype=np.float32)\n\n    for start in tqdm(range(0, N, batch_size), desc=\"ProtT5 Embedding\"):\n        end = min(start + batch_size, N)\n\n        ids = torch.tensor(ids_np[start:end], device=device)\n        mask = torch.tensor(mask_np[start:end], device=device)\n\n        with torch.no_grad():\n            with torch.cuda.amp.autocast(enabled=(device.type == \"cuda\")):\n                outputs = model(\n                    input_ids=ids,\n                    attention_mask=mask,\n                    output_hidden_states=True,\n                    return_dict=True\n                )\n\n                # Extract chosen layer\n                h = outputs.hidden_states[layer_idx]   # (B, L, 1024)\n\n                # Masked mean pooling\n                mask_f = mask.unsqueeze(-1)\n                pooled = (h * mask_f).sum(1) / mask_f.sum(1).clamp(min=1e-9)\n\n        embeddings[start:end] = pooled.float().cpu().numpy()\n\n        del ids, mask, h, pooled\n\n    return embeddings\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:43.549716Z","iopub.execute_input":"2025-12-15T13:02:43.549929Z","iopub.status.idle":"2025-12-15T13:02:43.556995Z","shell.execute_reply.started":"2025-12-15T13:02:43.549912Z","shell.execute_reply":"2025-12-15T13:02:43.556261Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# emb_t5_train = prot_t5_embed(\n#     \"Rostlab/prot_t5_xl_uniref50\",\n#     train_ids, train_mask,\n#     batch_size=4\n# )\n\n# np.save(\"/kaggle/working/train_prott5.npy\", emb_t5_train)\n# del emb_t5_train\n\nprint(\"Training prott5 completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:43.557836Z","iopub.execute_input":"2025-12-15T13:02:43.558177Z","iopub.status.idle":"2025-12-15T13:02:43.611268Z","shell.execute_reply.started":"2025-12-15T13:02:43.558158Z","shell.execute_reply":"2025-12-15T13:02:43.610638Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**test chunks were large so that's why I am chunking and combining them**","metadata":{}},{"cell_type":"code","source":"# MODEL_NAME = \"Rostlab/prot_t5_xl_uniref50\"\n# MAX_LEN = 512\n# LAYER_IDX = 12\n# EMB_DIM = 1024\n\n# MICRO_BATCH = 4                 # per TPU core (SAFE)\n# NUM_CORES = len(xm.get_xla_supported_devices())  # ✅ ALWAYS WORKS\n# GLOBAL_BATCH = MICRO_BATCH * NUM_CORES            # 32\n\n# CHUNK_SIZE = 4096               # proteins per chunk (~2–3 min)\n# OUT_DIR = \"/kaggle/working/prot_t5_test_chunks\"\n# os.makedirs(OUT_DIR, exist_ok=True)\n\n# print(\" TPU devices:\", xm.get_xla_supported_devices())\n# print(\" TPU cores:\", NUM_CORES)\n# print(\" Global batch:\", GLOBAL_BATCH)\n\n# # ======================================================\n# # DATA\n# # ======================================================\n# # test_ids, test_mask must already exist\n# N = test_ids.shape[0]\n# print(\"Total test proteins:\", N)\n\n# # ======================================================\n# # MODEL (TPU)\n# # ======================================================\n# device = xm.xla_device()\n\n# model = T5EncoderModel.from_pretrained(\n#     MODEL_NAME,\n#     torch_dtype=torch.bfloat16\n# ).to(device)\n\n# model.eval()\n\n# # ======================================================\n# # RESUME LOGIC\n# # ======================================================\n# def get_next_chunk_id(out_dir):\n#     files = [f for f in os.listdir(out_dir) if f.startswith(\"chunk_\")]\n#     if not files:\n#         return 0\n#     ids = sorted(int(f.split(\"_\")[1].split(\".\")[0]) for f in files)\n#     return ids[-1] + 1\n\n# start_chunk = get_next_chunk_id(OUT_DIR)\n# print(\" Resuming from chunk:\", start_chunk)\n\n# # ======================================================\n# # MAIN LOOP\n# # ======================================================\n# start_time = time.time()\n# num_chunks = math.ceil(N / CHUNK_SIZE)\n\n# for chunk_id in range(start_chunk, num_chunks):\n#     c_start = chunk_id * CHUNK_SIZE\n#     c_end = min(N, c_start + CHUNK_SIZE)\n#     cur_N = c_end - c_start\n\n#     print(f\"\\n Chunk {chunk_id} | proteins {c_start}:{c_end}\")\n\n#     out_chunk = np.zeros((cur_N, EMB_DIM), dtype=np.float32)\n#     write_ptr = 0\n\n#     for i in tqdm(\n#         range(c_start, c_end, GLOBAL_BATCH),\n#         desc=f\"TPU Chunk {chunk_id}\",\n#         leave=False\n#     ):\n#         j = min(i + GLOBAL_BATCH, c_end)\n#         bs = j - i\n\n#         ids = test_ids[i:j, :MAX_LEN]\n#         mask = test_mask[i:j, :MAX_LEN]\n\n#         # Pad to GLOBAL_BATCH\n#         if bs < GLOBAL_BATCH:\n#             pad = GLOBAL_BATCH - bs\n#             ids = np.pad(ids, ((0, pad), (0, 0)))\n#             mask = np.pad(mask, ((0, pad), (0, 0)))\n\n#         ids_t = torch.tensor(ids, device=device)\n#         mask_t = torch.tensor(mask, device=device)\n\n#         with torch.no_grad():\n#             out = model(\n#                 input_ids=ids_t,\n#                 attention_mask=mask_t,\n#                 output_hidden_states=True,\n#                 return_dict=True\n#             )\n\n#             h = out.hidden_states[LAYER_IDX]     # (B, L, 1024)\n#             mask_f = mask_t.unsqueeze(-1).to(h.dtype)\n#             pooled = (h * mask_f).sum(1) / mask_f.sum(1).clamp(min=1e-9)\n\n#         pooled = pooled[:bs].float().cpu().numpy()\n#         out_chunk[write_ptr:write_ptr + bs] = pooled\n#         write_ptr += bs\n\n#         xm.mark_step()   # TPU sync\n#         del ids_t, mask_t, h, pooled\n#         gc.collect()\n\n#     # SAVE CHUNK\n#     chunk_path = os.path.join(OUT_DIR, f\"chunk_{chunk_id:03d}.npy\")\n#     np.save(chunk_path, out_chunk)\n#     print(f\" Saved {chunk_path}\")\n\n# elapsed = (time.time() - start_time) / 60\nprint(f\"\\n DONE with ProtT5 Embeddings\")\n# print(f\" Total time: {elapsed:.1f} minutes\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:43.611905Z","iopub.execute_input":"2025-12-15T13:02:43.612856Z","iopub.status.idle":"2025-12-15T13:02:43.625980Z","shell.execute_reply.started":"2025-12-15T13:02:43.612839Z","shell.execute_reply":"2025-12-15T13:02:43.625455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# files = sorted(glob.glob(\"/kaggle/working/prot_t5_test_chunks/chunk_*.npy\"))\n# X = np.concatenate([np.load(f) for f in files], axis=0)\n\n# np.save(\"/kaggle/working/prot_t5_test.npy\", X)\nprint(\" Final test embeddings saved\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:43.626663Z","iopub.execute_input":"2025-12-15T13:02:43.626866Z","iopub.status.idle":"2025-12-15T13:02:43.645035Z","shell.execute_reply.started":"2025-12-15T13:02:43.626852Z","shell.execute_reply":"2025-12-15T13:02:43.644338Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  Part 4: Training the Single-Head Classifier\n\nWe will train a separate **Single-Layer Perceptron (MLP)** for each embedding type.  SLP would have three heads:  Molecular Function (MF), Biological Process (BP), Cellular Component (CC).\n\nI tried using Multi-Head classifier but the problem is I was constantly getting Memoryerror that's why swithced to single head","metadata":{}},{"cell_type":"code","source":"BASE_EMB = \"/kaggle/input/embeddings-for-cafa/\"\nBASE_LABELS = \"/kaggle/working/outputs_labels/\"\n\nfrom scipy import sparse\nimport numpy as np\n\nlabels_MF = sparse.load_npz(BASE_LABELS + \"labels_MF.npz\")\nlabels_BP = sparse.load_npz(BASE_LABELS + \"labels_BP.npz\")\nlabels_CC = sparse.load_npz(BASE_LABELS + \"labels_CC.npz\")\n\nlabels_by_ns = {\n    \"MF\": labels_MF,\n    \"BP\": labels_BP,\n    \"CC\": labels_CC\n}\n\nout_dims = {ns: mat.shape[1] for ns, mat in labels_by_ns.items()}\nprint(\"Loaded labels:\", out_dims)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:43.645798Z","iopub.execute_input":"2025-12-15T13:02:43.646122Z","iopub.status.idle":"2025-12-15T13:02:43.688184Z","shell.execute_reply.started":"2025-12-15T13:02:43.646097Z","shell.execute_reply":"2025-12-15T13:02:43.687608Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import train_test_split\n\n# All protein indices\nall_ids = np.arange(labels_MF.shape[0])\n\n# 10% validation split\ntrain_idx, val_idx = train_test_split(all_ids, test_size=0.1, random_state=42)\n\ndef split_labels(mat):\n    return mat[train_idx], mat[val_idx]\n\nlabels_train = {}\nlabels_val = {}\n\n# Split MF/BP/CC label matrices\nfor ns, mat in labels_by_ns.items():\n    labels_train[ns], labels_val[ns] = split_labels(mat)\n\n# Load embeddings\ntrain_emb_prot = np.load(BASE_EMB + \"train_protbert.npy\", mmap_mode=\"r\")\ntrain_emb_esm  = np.load(BASE_EMB + \"train_esm2.npy\", mmap_mode=\"r\")\n\n# Split embeddings\nprot_train = train_emb_prot[train_idx]\nprot_val   = train_emb_prot[val_idx]\n\nesm_train  = train_emb_esm[train_idx]\nesm_val    = train_emb_esm[val_idx]\n\nprint(\"Train/Val Split Complete\")\nprint(\"ProtBERT train:\", prot_train.shape, \"val:\", prot_val.shape)\nprint(\"ESM2 train:\", esm_train.shape, \"val:\", esm_val.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:43.688891Z","iopub.execute_input":"2025-12-15T13:02:43.689246Z","iopub.status.idle":"2025-12-15T13:02:57.257218Z","shell.execute_reply.started":"2025-12-15T13:02:43.689225Z","shell.execute_reply":"2025-12-15T13:02:57.256476Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\nUSE_TPU = False\ntry:\n    import torch_xla.core.xla_model as xm\n    device = xm.xla_device()\n    USE_TPU = True\n    print(\"TPU detected\")\nexcept Exception:\n    if torch.cuda.is_available():\n        device = torch.device(\"cuda\")\n        print(\"CUDA detected\")\n    else:\n        device = torch.device(\"cpu\")\n        print(\"CPU fallback\")\n\nprint(\"Using device:\", device)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:57.258084Z","iopub.execute_input":"2025-12-15T13:02:57.258336Z","iopub.status.idle":"2025-12-15T13:02:57.264715Z","shell.execute_reply.started":"2025-12-15T13:02:57.258307Z","shell.execute_reply":"2025-12-15T13:02:57.264080Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SingleHeadMLP(nn.Module):\n    def __init__(self, input_dim, proj_dim, out_dim, dropout):\n        super().__init__()\n\n        self.ln = nn.LayerNorm(input_dim)\n\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, proj_dim),\n            nn.LayerNorm(proj_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(proj_dim, proj_dim),\n            nn.LayerNorm(proj_dim),\n            nn.GELU(),\n            nn.Dropout(dropout),\n            nn.Linear(proj_dim, out_dim)\n        )\n\n    def forward(self, x):\n        x = self.ln(x)\n        return self.net(x)\n\nclass ProteinDatasetNS(Dataset):\n    def __init__(self, emb, labels):\n        self.emb = emb\n        self.labels = labels\n\n    def __len__(self):\n        return self.emb.shape[0]\n\n    def __getitem__(self, idx):\n        x = torch.tensor(self.emb[idx], dtype=torch.float32)\n        x = torch.nn.functional.normalize(x, dim=0)  # 🔥 FIX\n        y = torch.tensor(self.labels[idx].toarray().ravel(), dtype=torch.float32)\n        return x, y\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:57.265494Z","iopub.execute_input":"2025-12-15T13:02:57.265819Z","iopub.status.idle":"2025-12-15T13:02:57.281274Z","shell.execute_reply.started":"2025-12-15T13:02:57.265795Z","shell.execute_reply":"2025-12-15T13:02:57.280567Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"NS_CONFIG = {\n    \"MF\": dict(proj=512, dropout=0.2),\n    \"BP\": dict(proj=1024, dropout=0.3),\n    \"CC\": dict(proj=512, dropout=0.25),\n}\n\ndef train_ns_only(model_name, ns, emb_train, emb_val, y_train, y_val,\n                  batch_size=16, epochs=12, patience=2):\n\n    cfg = NS_CONFIG[ns]\n    print(f\"\\n[{model_name} | {ns}] Training\")\n\n    train_ds = ProteinDatasetNS(emb_train, y_train)\n    val_ds   = ProteinDatasetNS(emb_val, y_val)\n\n    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n    val_dl   = DataLoader(val_ds, batch_size=batch_size)\n\n    model = SingleHeadMLP(\n        input_dim=emb_train.shape[1],\n        proj_dim=cfg[\"proj\"],\n        out_dim=y_train.shape[1],\n        dropout=cfg[\"dropout\"]\n    ).to(device)\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n\n    label_counts = np.array(y_train.sum(axis=0)).ravel()\n    neg = y_train.shape[0] - label_counts\n    pos_weight = torch.tensor(np.clip(neg / (label_counts + 1e-6), 1, 50), device=device)\n\n    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n\n    best_val = float(\"inf\")\n    bad_epochs = 0\n\n    for ep in range(epochs):\n        model.train()\n        for x, y in tqdm(train_dl, leave=False):\n            x, y = x.to(device), y.to(device)\n            optimizer.zero_grad()\n            loss = criterion(model(x), y)\n            loss.backward()\n            if USE_TPU:\n                xm.optimizer_step(optimizer)\n                xm.mark_step()\n            else:\n                optimizer.step()\n\n        model.eval()\n        val_loss = 0.0\n        with torch.no_grad():\n            for x, y in val_dl:\n                x, y = x.to(device), y.to(device)\n                val_loss += criterion(model(x), y).item()\n\n        val_loss /= len(val_dl)\n        print(f\"Epoch {ep+1} | Val loss: {val_loss:.4f}\")\n\n        if val_loss < best_val:\n            best_val = val_loss\n            bad_epochs = 0\n        else:\n            bad_epochs += 1\n            if bad_epochs >= patience:\n                print(\"Early stopping\")\n                break\n\n        scheduler.step()\n\n    return model\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:57.282049Z","iopub.execute_input":"2025-12-15T13:02:57.282211Z","iopub.status.idle":"2025-12-15T13:02:57.295731Z","shell.execute_reply.started":"2025-12-15T13:02:57.282198Z","shell.execute_reply":"2025-12-15T13:02:57.295121Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"namespaces = [\"MF\", \"BP\", \"CC\"]\n\nprotbert_models = {}\nesm2_models = {}\nprott5_models = {}\n\n# Load ProtT5 embeddings (already generated)\ntrain_emb_t5 = np.load(BASE_EMB + \"train_prott5.npy\", mmap_mode=\"r\")\nval_emb_t5   = train_emb_t5[val_idx]\n\nfor ns in namespaces:\n    protbert_models[ns] = train_ns_only(\n        \"ProtBERT\", ns,\n        prot_train, prot_val,\n        labels_train[ns], labels_val[ns]\n    )\n\n    esm2_models[ns] = train_ns_only(\n        \"ESM2\", ns,\n        esm_train, esm_val,\n        labels_train[ns], labels_val[ns]\n    )\n\n    prott5_models[ns] = train_ns_only(\n        \"ProtT5\", ns,\n        train_emb_t5[train_idx], val_emb_t5,\n        labels_train[ns], labels_val[ns]\n    )\nprint(\"\\n Training Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:02:57.298380Z","iopub.execute_input":"2025-12-15T13:02:57.298651Z","iopub.status.idle":"2025-12-15T13:44:03.313698Z","shell.execute_reply.started":"2025-12-15T13:02:57.298636Z","shell.execute_reply":"2025-12-15T13:44:03.312873Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  Part 5: Submission\n\nAfter completing training for all namespaces (MF, BP, CC) with both the ProtBERT and ESM2-150M models, we generate predictions for the full CAFA-6 test set.\nInstead of saving massive .npy files (which corrupt easily and exceed Kaggle’s storage limits), we execute a streamed ensemble pipeline","metadata":{}},{"cell_type":"code","source":"with open(os.path.join(OUTPUT_LABELS_DIR, \"maps.json\")) as f:\n    maps = json.load(f)\nwith open(os.path.join(OUTPUT_LABELS_DIR, \"ancestors.json\")) as f:\n    ancestors = json.load(f)\n\nmaps_idx_to_go = {}\nfor ns in namespaces:\n    arr = np.empty(len(maps[ns]), dtype=object)\n    for i in range(len(arr)):\n        arr[i] = maps[ns][str(i)]\n    maps_idx_to_go[ns] = arr\n\nancestors_fast = {k: tuple(v) for k, v in ancestors.items()}\n\ntest_ids  = np.load(os.path.join(BASE_EMB, \"test_ids.npy\"))\ntest_prot = np.load(os.path.join(BASE_EMB, \"test_protbert.npy\"), mmap_mode=\"r\")\ntest_esm  = np.load(os.path.join(BASE_EMB, \"test_esm2.npy\"), mmap_mode=\"r\")\ntest_t5   = np.load(os.path.join(BASE_EMB, \"test_prott5.npy\"), mmap_mode=\"r\")\n\nN = len(test_ids)\nprint(\" Test proteins:\", N)\ntemperatures = {\"MF\": 0.85, \"BP\": 0.75, \"CC\": 0.95}\nthresholds   = {\"MF\": 8e-4, \"BP\": 1.2e-3, \"CC\": 4e-4}\n\nTOPK_RAW   = 2000\nTOPK_FINAL = 1500\nBATCH_SIZE = 256   \n\nsubmission_path = os.path.join(WORKING_DIR, \"submission.tsv\")\n\nwith open(submission_path, \"w\") as fout:\n    for start in tqdm(range(0, N, BATCH_SIZE), desc=\"Generating submission\"):\n        end = min(N, start + BATCH_SIZE)\n\n        prot = torch.as_tensor(test_prot[start:end], device=device)\n        esm  = torch.as_tensor(test_esm[start:end],  device=device)\n        t5   = torch.as_tensor(test_t5[start:end],   device=device)\n\n        batch_ids = test_ids[start:end]\n        ns_scores = {}\n\n        with torch.no_grad():\n            for ns in namespaces:\n                pb = torch.sigmoid(protbert_models[ns](prot))\n                es = torch.sigmoid(esm2_models[ns](esm))\n                tt = torch.sigmoid(prott5_models[ns](t5))\n\n                scores = torch.maximum(torch.maximum(pb, es), tt)\n                scores = scores / temperatures[ns]\n                ns_scores[ns] = scores.cpu().numpy()\n\n        for i, pid in enumerate(batch_ids):\n            scores_dict = {}\n\n            for ns in namespaces:\n                scores = ns_scores[ns][i]\n                th = thresholds[ns]\n\n                idxs = np.flatnonzero(scores > th)\n                if idxs.size == 0:\n                    continue\n\n                if idxs.size > TOPK_RAW:\n                    idxs = idxs[np.argpartition(scores[idxs], -TOPK_RAW)[-TOPK_RAW:]]\n\n                go_terms = maps_idx_to_go[ns][idxs]\n                vals = scores[idxs]\n\n                for go, sc in zip(go_terms, vals):\n                    prev = scores_dict.get(go)\n                    if prev is None or sc > prev:\n                        scores_dict[go] = sc\n\n                    for anc in ancestors_fast.get(go, ()):\n                        prev = scores_dict.get(anc)\n                        if prev is None or sc > prev:\n                            scores_dict[anc] = sc\n\n            if scores_dict:\n                items = sorted(scores_dict.items(), key=lambda x: x[1], reverse=True)[:TOPK_FINAL]\n                for go, sc in items:\n                    fout.write(f\"{pid}\\t{go}\\t{sc:.3f}\\n\")\n\n        del prot, esm, t5, ns_scores\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\nprint(\" submission.tsv written safely:\", submission_path)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-15T13:44:03.314602Z","iopub.execute_input":"2025-12-15T13:44:03.314876Z"}},"outputs":[],"execution_count":null}]}