{"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"},{"sourceId":5607816,"sourceType":"datasetVersion","datasetId":3225525},{"sourceId":14057766,"sourceType":"datasetVersion","datasetId":8947625},{"sourceId":14058756,"sourceType":"datasetVersion","datasetId":8948186},{"sourceId":517313,"sourceType":"modelInstanceVersion","modelInstanceId":407954,"modelId":425816}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# TSV FILES\nSAMPLE_SUBMISSION_TSV = \"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\"\nIA_TSV = \"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\"\nTESTSUPERSET_TAXON_LIST_TSV = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv\"\nTRAIN_TERMS_TSV = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\nTRAIN_TAXONOMY_TSV = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv\"\n\n# FASTA FILES\nTESTSUPERSET_FASTA = \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\nTRAIN_SEQUENCES_FASTA = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\n\n# OBO FILE\nGO_BASIC_OBO = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n\n# OUTPUT FILE\nOUTPUT_TSV = \"/kaggle/working/submission.tsv\"\n\nprint(\"Files are listed!!!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T22:54:06.166971Z","iopub.execute_input":"2025-12-08T22:54:06.167269Z","iopub.status.idle":"2025-12-08T22:54:06.176177Z","shell.execute_reply.started":"2025-12-08T22:54:06.167244Z","shell.execute_reply":"2025-12-08T22:54:06.17528Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# CONFIG\n# ------------------------------------------------------------\nCONFIG = {\n    \"TRAIN_FASTA\": TRAIN_SEQUENCES_FASTA,\n    \"TRAIN_TERMS\": TRAIN_TERMS_TSV,\n    \"TRAIN_TAXONOMY\": TRAIN_TAXONOMY_TSV,\n    \"GO_OBO\": GO_BASIC_OBO,\n    \"IA_FILE\": IA_TSV,\n    \"TEST_FASTA\": TESTSUPERSET_FASTA,\n    \"TEST_TAXON_LIST\": TESTSUPERSET_TAXON_LIST_TSV,\n    \"SAMPLE_SUBMISSION\": SAMPLE_SUBMISSION_TSV,\n    \"OUTPUT_SUBMISSION\": OUTPUT_TSV,\n    # Embedding settings:\n    \"USE_PLM_MODEL\": False,  # set False to force TF-IDF fallback\n    # If using TF/HF transformer model, either place checkpoint in dataset and point here,\n    # or use model name if internet enabled. On Kaggle usually you must provide local model.\n    \"PLM_MODEL_NAME_OR_PATH\": \"/kaggle/input/esm-2/keras/esm2_t6_8m/1\",  \n    \"PLM_BATCH_SIZE\": 8,\n    # Memory & batch sizes for streaming\n    \"EMBED_BATCH_SIZE\": 8,          # batch size used when embedding (train & test)\n    \"PREDICT_BATCH_SIZE\": 64,       # how many examples to predict at once (fit to memory)\n    # Label limitation to keep model small\n    \"TOP_K_LABELS\": 3000,\n    # Model training hyperparams\n    \"RANDOM_SEED\": 42,\n    \"BATCH_SIZE\": 32,\n    \"EPOCHS\": 10,\n    \"LEARNING_RATE\": 1e-3,\n    \"HIDDEN_UNITS\": 512,\n    \"DROPOUT\": 0.2,\n    # Submission postprocessing\n    \"TOP_K_PER_PROTEIN\": 200,\n    \"GLOBAL_THRESHOLD_SEARCH\": True,\n    \"THRESHOLD_GRID\": [i/100 for i in range(1, 51)],\n    # Propagation\n    \"PROPAGATE_TRAIN_LABELS\": True,\n    \"PROPAGATE_PREDICTIONS\": True,\n\n    # On-disk paths for memmaps/embeddings\n    \"TRAIN_EMB_MEMMAP\": \"/kaggle/working/train_embs.memmap\",\n    \"TRAIN_EMB_SHAPE_FILE\": \"/kaggle/working/train_embs_shape.npy\",\n}\n\n\nprint(\"config...\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T22:54:06.646357Z","iopub.execute_input":"2025-12-08T22:54:06.646653Z","iopub.status.idle":"2025-12-08T22:54:06.652473Z","shell.execute_reply.started":"2025-12-08T22:54:06.646634Z","shell.execute_reply":"2025-12-08T22:54:06.651754Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pip install -q git+https://github.com/facebookresearch/esm.git","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T22:54:06.653861Z","iopub.execute_input":"2025-12-08T22:54:06.654133Z","iopub.status.idle":"2025-12-08T22:54:19.78849Z","shell.execute_reply.started":"2025-12-08T22:54:06.654109Z","shell.execute_reply":"2025-12-08T22:54:19.787581Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# imports\n# ------------------------------------------------------------\nimport os, gc, math, random, sys, time\nfrom collections import defaultdict, Counter\nfrom typing import List, Dict, Tuple, Set\nimport numpy as np\nimport pandas as pd\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nimport tensorflow as tf\nfrom tensorflow.keras import layers, models, callbacks, optimizers, losses, metrics\nimport esm   # after installing from GitHub\nimport torch\nimport numpy as np\nimport gc\nfrom pathlib import Path\n\n# optional PLM imports\ntry:\n    import torch\n    from transformers import AutoTokenizer, AutoModel\n    TORCH_AVAILABLE = True\nexcept Exception as e:\n    TORCH_AVAILABLE = False\n\n# deterministic seeds\nrandom.seed(CONFIG[\"RANDOM_SEED\"])\nnp.random.seed(CONFIG[\"RANDOM_SEED\"])\ntf.random.set_seed(CONFIG[\"RANDOM_SEED\"])\n\nprint(\"import done!!!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T22:54:19.790648Z","iopub.execute_input":"2025-12-08T22:54:19.790933Z","iopub.status.idle":"2025-12-08T22:54:57.339169Z","shell.execute_reply.started":"2025-12-08T22:54:19.790908Z","shell.execute_reply":"2025-12-08T22:54:57.338485Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# tiny utils for FASTA/TSV/OBO parsing \n# ------------------------------------------------------------\ndef read_fasta(path: str) -> Dict[str, str]:\n    seqs = {}\n    with open(path, \"r\") as f:\n        pid = None; seq_parts = []\n        for line in f:\n            line=line.strip()\n            if line.startswith(\">\"):\n                if pid: seqs[pid] = \"\".join(seq_parts)\n                header=line[1:].split()[0]\n                if \"|\" in header:\n                    parts=header.split(\"|\"); pid = parts[1] if len(parts)>=2 else header\n                else:\n                    pid = header\n                seq_parts=[]\n            else:\n                seq_parts.append(line.strip())\n        if pid: seqs[pid] = \"\".join(seq_parts)\n    print(f\"[io] Read {len(seqs)} sequences from {path}\")\n    return seqs\n\ndef read_train_terms(path: str) -> Dict[str, List[str]]:\n    mapping = defaultdict(list)\n    df = pd.read_csv(path, sep=\"\\t\", header=None, names=[\"protein\",\"go\",\"ont\"], dtype=str)\n    for _, r in df.iterrows(): mapping[r.protein].append(r.go)\n    print(f\"[io] Read training annotations for {len(mapping)} proteins from {path}\")\n    return mapping\n\ndef parse_obo(go_obo_path: str) -> Tuple[Dict[str, Set[str]], Dict[str, Set[str]]]:\n    parents = defaultdict(set); children = defaultdict(set)\n    if not os.path.exists(go_obo_path): return parents, children\n    with open(go_obo_path,\"r\") as f:\n        cur_id=None\n        for line in f:\n            line=line.strip()\n            if line==\"[Term]\": cur_id=None\n            elif line.startswith(\"id: \"): cur_id=line.split(\"id: \")[1].strip()\n            elif line.startswith(\"is_a: \"):\n                pid=line.split()[1].strip()\n                if cur_id: parents[cur_id].add(pid); children[pid].add(cur_id)\n            elif line.startswith(\"relationship: part_of \"):\n                parts=line.split(); \n                if len(parts)>=3:\n                    pid=parts[2].strip()\n                    if cur_id: parents[cur_id].add(pid); children[pid].add(cur_id)\n    print(f\"[io] Parsed OBO: {len(parents)} nodes with parents\")\n    return parents, children\n\ndef get_ancestors(go_id: str, parents: Dict[str, Set[str]]) -> Set[str]:\n    ans=set(); stack=[go_id]\n    while stack:\n        cur=stack.pop()\n        for p in parents.get(cur,[]): \n            if p not in ans:\n                ans.add(p); stack.append(p)\n    return ans\ndef read_IA(path: str) -> pd.DataFrame:\n    # đọc file TSV, mỗi dòng gồm go và score\n    df = pd.read_csv(path, sep=\"\\t\", header=None, names=[\"go\", \"score\"])\n    df[\"score\"] = df[\"score\"].astype(float)  # chắc chắn score là float\n    return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T22:54:57.339883Z","iopub.execute_input":"2025-12-08T22:54:57.3403Z","iopub.status.idle":"2025-12-08T22:54:57.351496Z","shell.execute_reply.started":"2025-12-08T22:54:57.340282Z","shell.execute_reply":"2025-12-08T22:54:57.35073Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Load the data\n# ------------------------------------------------------------\ntrain_seqs = read_fasta(CONFIG[\"TRAIN_FASTA\"])\ntrain_terms = read_train_terms(CONFIG[\"TRAIN_TERMS\"])\nparents_map, children_map = parse_obo(CONFIG[\"GO_OBO\"])\ntest_seqs = read_fasta(CONFIG[\"TEST_FASTA\"])\n\n# Keep proteins present in both seq & terms\ntrain_proteins = [p for p in train_terms.keys() if p in train_seqs]\nprint(f\"[io] {len(train_proteins)} train proteins with sequences available\")\n\n# propagate train labels (optional)\nif CONFIG[\"PROPAGATE_TRAIN_LABELS\"] and parents_map:\n    print(\"[prep] Propagating train labels up GO graph\")\n    propagated={}\n    for p in train_proteins:\n        terms=set(train_terms[p])\n        extra=set()\n        for t in list(terms): extra |= get_ancestors(t, parents_map)\n        propagated[p]=sorted(terms|extra)\n    train_terms = propagated\n\n# choose top-k labels (to control model size)\nall_term_counts = Counter()\nfor p in train_proteins: all_term_counts.update(train_terms[p])\nall_terms_sorted = [t for t,_ in all_term_counts.most_common()]\nif CONFIG[\"TOP_K_LABELS\"] is not None:\n    chosen_terms = set(all_terms_sorted[:CONFIG[\"TOP_K_LABELS\"]])\n    print(f\"[prep] Restricting to top-{CONFIG['TOP_K_LABELS']} GO terms\")\nelse:\n    chosen_terms = set(all_terms_sorted)\nprint(f\"[prep] Using {len(chosen_terms)} target GO terms\")\n\nfor p in train_proteins:\n    train_terms[p] = [t for t in train_terms[p] if t in chosen_terms]\n\nX_proteins = train_proteins\ny_labels = [train_terms[p] for p in X_proteins]\n\nmlb = MultiLabelBinarizer(classes=sorted(chosen_terms))\nY = mlb.fit_transform(y_labels).astype(np.float32)\nprint(\"[prep] Label matrix shape:\", Y.shape)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Load the data\n# ------------------------------------------------------------\ntrain_seqs = read_fasta(CONFIG[\"TRAIN_FASTA\"])\ntrain_terms = read_train_terms(CONFIG[\"TRAIN_TERMS\"])\nparents_map, children_map = parse_obo(CONFIG[\"GO_OBO\"])\ntest_seqs = read_fasta(CONFIG[\"TEST_FASTA\"])\n\n# Keep proteins present in both seq & terms\ntrain_proteins = [p for p in train_terms.keys() if p in train_seqs]\n\nprint(f\"[io] {len(train_proteins)} train proteins with sequences available\")\n\nIA_list = read_IA(CONFIG[\"IA_FILE\"])\n\n# propagate train labels (optional)\nif CONFIG[\"PROPAGATE_TRAIN_LABELS\"] and parents_map:\n    print(\"[prep] Propagating train labels up GO graph\")\n    propagated={}\n    for p in train_proteins:\n        terms=set(train_terms[p])\n        extra=set()\n        for t in list(terms): extra |= get_ancestors(t, parents_map)\n        propagated[p]=sorted(terms|extra)\n    train_terms = propagated\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T22:59:33.840078Z","iopub.execute_input":"2025-12-08T22:59:33.840838Z","iopub.status.idle":"2025-12-08T23:03:47.94282Z","shell.execute_reply.started":"2025-12-08T22:59:33.840798Z","shell.execute_reply":"2025-12-08T23:03:47.942192Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\n\n# số GO term của mỗi protein\ngo_counts = np.array([len(train_terms[p]) for p in train_proteins])\n\nprint(\"========== GO TERM DISTRIBUTION ==========\")\nprint(\"Tổng protein:\", len(go_counts))\nprint(\"Min GO / protein:\", go_counts.min())\nprint(\"Max GO / protein:\", go_counts.max())\nprint(\"Mean GO / protein:\", go_counts.mean())\nprint(\"Median GO / protein:\", np.median(go_counts))\nprint(\"Std GO / protein:\", go_counts.std())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:15:18.764938Z","iopub.execute_input":"2025-12-08T23:15:18.765574Z","iopub.status.idle":"2025-12-08T23:15:18.844232Z","shell.execute_reply.started":"2025-12-08T23:15:18.765549Z","shell.execute_reply":"2025-12-08T23:15:18.843419Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"IA_list = read_IA(CONFIG[\"IA_FILE\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:04:11.615878Z","iopub.execute_input":"2025-12-08T23:04:11.616192Z","iopub.status.idle":"2025-12-08T23:04:11.639943Z","shell.execute_reply.started":"2025-12-08T23:04:11.616168Z","shell.execute_reply":"2025-12-08T23:04:11.639192Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(type(IA_list))\nprint(IA_list.head())\nprint(IA_list.columns)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:04:13.722613Z","iopub.execute_input":"2025-12-08T23:04:13.72321Z","iopub.status.idle":"2025-12-08T23:04:13.742958Z","shell.execute_reply.started":"2025-12-08T23:04:13.723184Z","shell.execute_reply":"2025-12-08T23:04:13.742236Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"len(IA_list)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:04:16.447366Z","iopub.execute_input":"2025-12-08T23:04:16.448044Z","iopub.status.idle":"2025-12-08T23:04:16.453613Z","shell.execute_reply.started":"2025-12-08T23:04:16.44802Z","shell.execute_reply":"2025-12-08T23:04:16.452881Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- đọc IA file: IA_list = {go_id: score} ---\nfrom collections import Counter\nimport numpy as np\n\n# --- tính term frequency ---\nterm_freq = Counter()\nfor p in train_proteins:\n    term_freq.update(train_terms[p])   # giữ nguyên \"GO:xxxx\"\n\n# --- tạo dict tra nhanh từ DataFrame (không chuẩn hóa) ---\nIA_dict = dict(zip(IA_list[\"go\"], IA_list[\"score\"]))\n\n# --- tính combined score ---\nterm_combined_score = {}\nfor t in term_freq:\n    ia_score = IA_dict.get(t, 0)   # giữ nguyên GO:xxxx\n    term_combined_score[t] = ia_score * term_freq[t]\n\n# --- chọn GO từ hạng 3001 -> 6000 theo combined score ---\nRANK_START = 3000   # bỏ top 3000 đầu\nRANK_END   = 6000   # lấy tới 6000\n\n# --- chọn top-K theo combined score ---\nif CONFIG[\"TOP_K_LABELS\"] is not None:\n    top_terms_sorted = sorted(term_combined_score.items(), key=lambda x: x[1], reverse=True)\n    chosen_terms = set([t for t, _ in top_terms_sorted[:CONFIG[\"TOP_K_LABELS\"]]])\n    print(f\"[prep] Restricting to top-{CONFIG['TOP_K_LABELS']} GO terms by IA*freq\")\n    #chosen_terms = set([t for t, _ in top_terms_sorted[RANK_START:RANK_END]])\n    #print(f\"[prep] Using GO ranks {RANK_START+1} -> {RANK_END}\")\nelse:\n    chosen_terms = set(term_combined_score.keys())\nprint(f\"[prep] Using {len(chosen_terms)} target GO terms\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:04:18.042945Z","iopub.execute_input":"2025-12-08T23:04:18.043614Z","iopub.status.idle":"2025-12-08T23:04:18.624221Z","shell.execute_reply.started":"2025-12-08T23:04:18.043586Z","shell.execute_reply":"2025-12-08T23:04:18.623467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\n# --- sắp xếp theo combined score giảm dần ---\nsorted_scores = sorted(term_combined_score.items(), key=lambda x: x[1], reverse=True)\nscores = np.array([s for _, s in sorted_scores])\n\n# --- tính phân phối tích lũy ---\ncum_scores = np.cumsum(scores)\ntotal_score = cum_scores[-1]\ncum_ratio = cum_scores / total_score\n\n# --- tìm K sao cho đạt 90% ---\nK_90 = np.searchsorted(cum_ratio, 0.9) + 1\nK_95 = np.searchsorted(cum_ratio, 0.95) + 1\nK_99 = np.searchsorted(cum_ratio, 0.99) + 1\n\nprint(f\"[dist] Total GO terms: {len(scores)}\")\nprint(f\"[dist] Top-{K_90} GO terms cover 90% combined score\")\nprint(f\"[dist] Top-{K_95} GO terms cover 95% combined score\")\nprint(f\"[dist] Top-{K_99} GO terms cover 99% combined score\")\n","metadata":{"execution":{"iopub.status.busy":"2025-12-07T06:55:37.199424Z","iopub.execute_input":"2025-12-07T06:55:37.199978Z","iopub.status.idle":"2025-12-07T06:55:37.209576Z","shell.execute_reply.started":"2025-12-07T06:55:37.199951Z","shell.execute_reply":"2025-12-07T06:55:37.20892Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"total_score = sum(term_combined_score.values())\nnon_zero = sum(1 for v in term_combined_score.values() if v > 0)\n\nprint(\"Total combined score:\", total_score)\nprint(\"Non-zero GO terms:\", non_zero)\n\n# Xem 10 GO term lớn nhất\ntop10 = sorted(term_combined_score.items(), key=lambda x: x[1], reverse=True)[:10]\nfor t, s in top10:\n    print(t, s)","metadata":{"execution":{"iopub.status.busy":"2025-12-07T06:56:23.781381Z","iopub.execute_input":"2025-12-07T06:56:23.782053Z","iopub.status.idle":"2025-12-07T06:56:23.786898Z","shell.execute_reply.started":"2025-12-07T06:56:23.782028Z","shell.execute_reply":"2025-12-07T06:56:23.786058Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\nplt.figure()\nplt.plot(cum_ratio)\nplt.axhline(0.9)\nplt.axhline(0.95)\nplt.axhline(0.99)\nplt.xlabel(\"Top-K GO terms\")\nplt.ylabel(\"Cumulative ratio of IA*freq\")\nplt.title(\"Coverage of GO terms by IA*freq\")\nplt.grid(True)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2025-12-07T06:56:26.969528Z","iopub.execute_input":"2025-12-07T06:56:26.970174Z","iopub.status.idle":"2025-12-07T06:56:27.114649Z","shell.execute_reply.started":"2025-12-07T06:56:26.970151Z","shell.execute_reply":"2025-12-07T06:56:27.11398Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# PLM embedding helpers (ESM2); memory-conscious: produce numpy arrays per-batch\n# ------------------------------------------------------------\ndef seqs_for_plm_input_esm(seqs: List[str]) -> List[str]:\n    # ESM expects raw sequences (no spaces); we simply uppercase and replace unknowns\n    out=[]\n    for s in seqs:\n        s2 = s.upper().replace(\"U\",\"X\").replace(\"O\",\"X\").replace(\"B\",\"X\").replace(\"Z\",\"X\")\n        out.append(s2)\n    return out\n\ndef embed_with_plm_to_memmap(all_seq_ids: List[str],\n                             seqs_dict: Dict[str,str],\n                             memmap_path: str,\n                             batch_size:int=8,\n                             model_name_or_path: str = CONFIG[\"PLM_MODEL_NAME_OR_PATH\"]):\n    \"\"\"\n    Compute embeddings using the ESM loader and write to disk-backed memmap.\n    Returns memmap object and embedding dimension.\n    This function expects the local directory model_name_or_path to contain the ESM checkpoint\n    produced by the ESM tooling (e.g., esm2_t33_650M_UR50D.pt or a model dir).\n    \"\"\"\n    if not TORCH_AVAILABLE:\n        raise RuntimeError(\"Torch not available; cannot load ESM model.\")\n\n    model_dir = str(model_name_or_path)\n    if not Path(model_dir).exists():\n        raise FileNotFoundError(f\"ESM model path not found: {model_dir}\")\n\n    # Try to load via esm loader that understands local formats\n    try:\n        # If model_name_or_path is a directory that contains a model checkpoint,\n        # this loader will try to read it. If it points directly to a .pt file it also works.\n        print(f\"[esm] Loading local ESM model from: {model_dir}\")\n        model, alphabet = esm.pretrained.load_model_and_alphabet_local(model_dir)\n    except Exception as e:\n        # some ESM checkpoints use slightly different utilities - try the convenience function names:\n        try:\n            # If the user packaged a directory like \"esm2_t33_650M_UR50D\" that contains model.pt\n            print(f\"[esm] load_model_and_alphabet_local failed, attempting esm2 convenience loader...\")\n            model, alphabet = esm.pretrained.esm2_t6_8M_UR50D()  # fallback: only works if model files are accessible\n        except Exception as e2:\n            raise RuntimeError(\"Failed to load ESM model via esm.pretrained. Ensure the directory contains a valid ESM checkpoint.\") from e\n\n    model.eval()\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model.to(device)\n\n    batch_converter = alphabet.get_batch_converter()\n    N = len(all_seq_ids)\n\n    # Determine embedding dimension using a tiny sample (1 sequence)\n    sample_seq = seqs_dict[all_seq_ids[0]]\n    _, _, sample_tokens = batch_converter([(all_seq_ids[0], sample_seq)])\n    sample_tokens = sample_tokens.to(device)\n    with torch.no_grad():\n        results = model(sample_tokens, repr_layers=[model.num_layers], return_contacts=False)\n        # pick the highest repr layer in results\n        repr_keys = sorted(results[\"representations\"].keys())\n        last_layer_key = repr_keys[-1]\n        emb_dim = results[\"representations\"][last_layer_key].shape[-1]\n    # create memmap file on disk (float32)\n    mem = np.memmap(memmap_path, dtype=np.float32, mode=\"w+\", shape=(N, int(emb_dim)))\n\n    idx = 0\n    for i in range(0, N, batch_size):\n        batch_ids = all_seq_ids[i:i+batch_size]\n        # Prepare list of (id, seq) tuples\n        batch_pairs = [(pid, seqs_dict[pid]) for pid in batch_ids]\n        labels, sequences, tokens = batch_converter(batch_pairs)  # tokens: (B, L)\n        tokens = tokens.to(device)\n        with torch.no_grad():\n            results = model(tokens, repr_layers=[model.num_layers], return_contacts=False)\n            repr_keys = sorted(results[\"representations\"].keys())\n            last_layer_key = repr_keys[-1]\n            repr_tensor = results[\"representations\"][last_layer_key].cpu()   # (B, L, C)\n        # For each sequence, slice out residues and mean-pool (drop BOS/EOS token at positions 0 and -1)\n        for j, seq in enumerate(sequences):\n            seq_len = len(seq)\n            # ESM token layout: tokens include BOS at pos 0 and EOS at pos seq_len+1 so we slice 1:seq_len+1\n            seq_repr = repr_tensor[j, 1:seq_len+1, :]   # (seq_len, C)\n            seq_embed = seq_repr.mean(axis=0).numpy().astype(np.float32)\n            mem[idx + j, :] = seq_embed\n        idx += len(batch_ids)\n\n        # free intermediate tensors and empty CUDA cache\n        del tokens, results, repr_tensor, seq_repr, seq_embed\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n        print(f\"[esm] Wrote embeddings {i}..{i+len(batch_ids)} / {N}\")\n\n    mem.flush()\n    print(f\"[esm] Finished writing memmap to {memmap_path} with dim {emb_dim}\")\n    return mem, int(emb_dim)","metadata":{"execution":{"execution_failed":"2025-12-04T00:39:11.698Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Embedding training set (memmap) OR fallback TF-IDF (also memory-friendly)\n# ------------------------------------------------------------\nUSE_PLM = CONFIG[\"USE_PLM_MODEL\"] and TORCH_AVAILABLE\nif USE_PLM:\n    # compute training embeddings to disk memmap to avoid storing huge array in RAM\n    train_ids = X_proteins\n    train_memmap_path = CONFIG[\"TRAIN_EMB_MEMMAP\"]\n    if not os.path.exists(train_memmap_path):\n        print(\"[plm] Computing train embeddings to memmap. This may take time but keeps RAM low.\")\n        train_mem, D = embed_with_plm_to_memmap(train_ids, train_seqs, train_memmap_path,\n                                                batch_size=CONFIG[\"EMBED_BATCH_SIZE\"],\n                                                model_name_or_path=CONFIG[\"PLM_MODEL_NAME_OR_PATH\"])\n        # Save shape info for later reopening\n        np.save(CONFIG[\"TRAIN_EMB_SHAPE_FILE\"], np.array([len(train_ids), D], dtype=np.int64))\n        # Keep mem as memmap object reference\n        emb_train = np.array(train_mem)  # small temporary conversion for training. If too large, we will reopen memmap later.\n        # To be safe, copy to np.float32 if necessary\n        if emb_train.dtype != np.float32: emb_train = emb_train.astype(np.float32)\n        del train_mem; gc.collect()\n        if torch.cuda.is_available(): torch.cuda.empty_cache()\n    else:\n        # memmap already exists: load shape and open\n        shape = np.load(CONFIG[\"TRAIN_EMB_SHAPE_FILE\"])\n        Nshape, D = int(shape[0]), int(shape[1])\n        emb_train = np.memmap(train_memmap_path, dtype=np.float32, mode=\"r\", shape=(Nshape, D))\n    embedding_method = \"plm\"\nelse:\n    # TF-IDF fallback (fits in memory usually; if not, you can also memmap)\n    print(\"[fallback] Using TF-IDF k-mer embeddings for train (memory-friendly for moderate sizes)\")\n    def kmers(seq, k=3):\n        return \" \".join([seq[i:i+k] for i in range(len(seq)-k+1)])\n    train_texts = [kmers(train_seqs[p], k=3) for p in X_proteins]\n    tfidf = TfidfVectorizer(analyzer=\"word\", token_pattern=r\"(?u)\\b\\w+\\b\", max_features=20000)\n    emb_train = tfidf.fit_transform(train_texts).astype(np.float32).toarray()\n    embedding_method = \"tfidf\"\n\nprint(f\"[embed] Train embeddings: method={embedding_method}, shape={emb_train.shape}, dtype={emb_train.dtype}\")\n","metadata":{"execution":{"execution_failed":"2025-12-04T00:39:11.698Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Train / validation split (we load train embeddings into memory here - if it's too big, we can train with memmap directly)\n# ------------------------------------------------------------\n# If embeddings are memmap and too big, we can load a subset or use a generator; here we assume emb_train fits for training.\nX_emb = emb_train\ny = Y\nX_tr, X_val, y_tr, y_val, prot_tr, prot_val = train_test_split(\n    X_emb, y, X_proteins, test_size=0.15, random_state=CONFIG[\"RANDOM_SEED\"]\n)\nprint(\"[train] shapes:\", X_tr.shape, X_val.shape, y_tr.shape, y_val.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:04:37.777209Z","iopub.execute_input":"2025-12-08T23:04:37.777996Z","iopub.status.idle":"2025-12-08T23:04:46.073112Z","shell.execute_reply.started":"2025-12-08T23:04:37.777969Z","shell.execute_reply":"2025-12-08T23:04:46.072411Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Keras model\n# ------------------------------------------------------------\nD = X_tr.shape[1]; M = y_tr.shape[1]\ndef build_model(input_dim, output_dim, hidden_units=512, dropout=0.4, lr=1e-3):\n    inp = layers.Input(shape=(input_dim,))\n    x = layers.Dense(hidden_units, activation=\"relu\")(inp)\n    x = layers.BatchNormalization()(x)\n    x = layers.Dropout(dropout)(x)\n    x = layers.Dense(hidden_units//2, activation=\"relu\")(x)\n    x = layers.BatchNormalization()(x)\n    x = layers.Dropout(dropout)(x)\n    out = layers.Dense(output_dim, activation=\"sigmoid\")(x)\n    model = models.Model(inp, out)\n    model.compile(optimizer=optimizers.Adam(learning_rate=lr),\n                  loss=losses.BinaryCrossentropy(),\n                  metrics=[metrics.Precision(), metrics.Recall()])\n    return model\n\nmodel = build_model(D, M, hidden_units=CONFIG[\"HIDDEN_UNITS\"], dropout=CONFIG[\"DROPOUT\"], lr=CONFIG[\"LEARNING_RATE\"])\nmodel.summary()\n\nes = callbacks.EarlyStopping(monitor=\"val_loss\", patience=3, restore_best_weights=True, verbose=1)\nckpt_path = \"/kaggle/working/best_model.h5\"\nmc = callbacks.ModelCheckpoint(ckpt_path, monitor=\"val_loss\", save_best_only=True, verbose=1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:18:33.390139Z","iopub.execute_input":"2025-12-09T18:18:33.390361Z","iopub.status.idle":"2025-12-09T18:18:33.465924Z","shell.execute_reply.started":"2025-12-09T18:18:33.390335Z","shell.execute_reply":"2025-12-09T18:18:33.465062Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Train\n# ------------------------------------------------------------\nhistory = model.fit(X_tr, y_tr, validation_data=(X_val, y_val),\n                    epochs=CONFIG[\"EPOCHS\"], batch_size=CONFIG[\"BATCH_SIZE\"],\n                    callbacks=[es, mc], verbose=2)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:05:16.647468Z","iopub.execute_input":"2025-12-08T23:05:16.648145Z","iopub.status.idle":"2025-12-08T23:07:10.287629Z","shell.execute_reply.started":"2025-12-08T23:05:16.648124Z","shell.execute_reply":"2025-12-08T23:07:10.28702Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Evaluate & select global threshold\n# ------------------------------------------------------------\ndef weighted_precision_recall_f1(y_true, y_pred_bin, ia_map, mlb_obj):\n    tp = ((y_true==1)&(y_pred_bin==1)).sum(axis=0).astype(float)\n    fp = ((y_true==0)&(y_pred_bin==1)).sum(axis=0).astype(float)\n    fn = ((y_true==1)&(y_pred_bin==0)).sum(axis=0).astype(float)\n    eps=1e-12\n    prec = tp/(tp+fp+eps); rec = tp/(tp+fn+eps)\n    f1 = 2*prec*rec/(prec+rec+eps)\n    cls = mlb_obj.classes_\n    weights = np.array([ia_weights.get(c,1.0) for c in cls], dtype=float) if 'ia_weights' in globals() else np.ones(len(cls))\n    weighted_f1 = (f1*weights).sum()/(weights.sum()+eps)\n    weighted_prec = (prec*weights).sum()/(weights.sum()+eps)\n    weighted_rec = (rec*weights).sum()/(weights.sum()+eps)\n    return weighted_prec, weighted_rec, weighted_f1\n\n# load IA weights if available (safe)\ndef read_IA_safe(path):\n    if not os.path.exists(path): return {}\n    df=pd.read_csv(path, sep=\"\\t\", header=None, names=[\"go\",\"ia\"], dtype=str)\n    d={}\n    for _,r in df.iterrows():\n        try: d[r.go]=float(r.ia)\n        except: \n            try: d[r.go]=float(r.ia.replace(\",\",\".\")) \n            except: d[r.go]=0.0\n    return d\n\nia_weights = read_IA_safe(CONFIG[\"IA_FILE\"])\n\ny_val_prob = model.predict(X_val, batch_size=CONFIG[\"BATCH_SIZE\"], verbose=0)\nbest_thresh = 0.5; best_score = -1.0\nif CONFIG[\"GLOBAL_THRESHOLD_SEARCH\"]:\n    for t in CONFIG[\"THRESHOLD_GRID\"]:\n        y_pred_bin = (y_val_prob >= t).astype(int)\n        _,_,f1 = weighted_precision_recall_f1(y_val, y_pred_bin, ia_weights, mlb)\n        if f1 > best_score: best_score=f1; best_thresh=t\nprint(f\"[eval] best_thresh {best_thresh} best IA-weighted F1 {best_score}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:07:16.997662Z","iopub.execute_input":"2025-12-08T23:07:16.998508Z","iopub.status.idle":"2025-12-08T23:07:24.915331Z","shell.execute_reply.started":"2025-12-08T23:07:16.998474Z","shell.execute_reply":"2025-12-08T23:07:24.91441Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ------------------------------------------------------------\n# Streaming test-time: embed test sequences batchwise, predict, propagate, write submission lines immediately\n# ------------------------------------------------------------\n# Precompute helpful mappings for propagation\nterm_to_idx = {t:i for i,t in enumerate(mlb.classes_)}\nidx_to_term = {i:t for t,i in term_to_idx.items()}\n\n# Build parents_map restricted to chosen_terms to speed propagation\nrestricted_parents = {}\nfor t in mlb.classes_:\n    restricted_parents[t] = set([p for p in parents_map.get(t, set()) if p in term_to_idx])\n\n# small propagation routine operating on a batch_of_probs (N_batch, M)\ndef propagate_batch(pred_batch: np.ndarray, parents_map_local: Dict[str, Set[str]], classes_list: List[str], iterations=3):\n    # pred_batch: float32 shape (B, M)\n    B, Mloc = pred_batch.shape\n    idx_map = {i:classes_list[i] for i in range(Mloc)}\n    term_to_idx_local = {classes_list[i]: i for i in range(Mloc)}\n    for _ in range(iterations):\n        changed = False\n        # vectorized-ish: for each child index, update parent index with max\n        # loop over terms (M might be a few thousands => ok per small batch)\n        for child_idx in range(Mloc):\n            child_term = idx_map[child_idx]\n            child_scores = pred_batch[:, child_idx]\n            for pterm in parents_map_local.get(child_term, []):\n                pidx = term_to_idx_local[pterm]\n                # update parent where child's score exceeds parent\n                mask = child_scores > pred_batch[:, pidx]\n                if mask.any():\n                    pred_batch[mask, pidx] = child_scores[mask]\n                    changed = True\n        if not changed: break\n    return pred_batch\n\n# Open submission file for streaming write\nout_fpath = CONFIG[\"OUTPUT_SUBMISSION\"]\nopen(out_fpath, \"w\").close()  # truncate\nout_f = open(out_fpath, \"a\")\n\n# Create chunked iterator of test sequence IDs\ntest_ids = list(test_seqs.keys())\nN_test = len(test_ids)\nprint(f\"[test] Streaming {N_test} test sequences in batches of {CONFIG['EMBED_BATCH_SIZE']} (embed) / predict {CONFIG['PREDICT_BATCH_SIZE']}\")\n\n# If using PLM, prepare tokenizer & model once (on CPU/GPU)\nif USE_PLM:\n    tokenizer = AutoTokenizer.from_pretrained(CONFIG[\"PLM_MODEL_NAME_OR_PATH\"], do_lower_case=False)\n    plm_model = AutoModel.from_pretrained(CONFIG[\"PLM_MODEL_NAME_OR_PATH\"])\n    plm_model.eval()\n    if torch.cuda.is_available(): plm_model.to(torch.device(\"cuda\"))\n\n# Helper to embed a list of sequences and return numpy array float32 of shape (len(seq_list), D)\ndef embed_batch_return_np(seq_list: List[str]):\n    if USE_PLM:\n        proc = seqs_for_plm_input_esm(seq_list)\n        with torch.no_grad():\n            inputs = tokenizer(proc, return_tensors=\"pt\", padding=True, truncation=True)\n            device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n            inputs = {k:v.to(device) for k,v in inputs.items()}\n            out = plm_model(**inputs)\n            last_hidden = out.last_hidden_state  # (B, L, dim)\n            mask = inputs.get(\"attention_mask\", None)\n            if mask is not None:\n                mask = mask.unsqueeze(-1)\n                summed = (last_hidden * mask).sum(1)\n                counts = mask.sum(1).clamp(min=1)\n                mean_pooled = (summed / counts).cpu().numpy().astype(np.float32)\n            else:\n                mean_pooled = last_hidden.mean(dim=1).cpu().numpy().astype(np.float32)\n        # free GPU memory\n        del inputs, out, last_hidden\n        if torch.cuda.is_available(): torch.cuda.empty_cache()\n        gc.collect()\n        return mean_pooled\n    else:\n        # TF-IDF fallback: transform using vectorizer in small chunk\n        texts = [\" \".join([seq[i:i+3] for i in range(len(seq)-3+1)]) for seq in seq_list]\n        arr = tfidf.transform(texts).astype(np.float32).toarray()\n        return arr\n\n# We'll process test samples in small \"embed\" batches and then accumulate a list of embedded rows\n# until we have PREDICT_BATCH_SIZE embeddings ready for model.predict(...) then predict and write submission lines.\nembed_batch = []\nembed_ids = []\n\nfor i in range(0, N_test, CONFIG[\"EMBED_BATCH_SIZE\"]):\n    batch_ids = test_ids[i:i+CONFIG[\"EMBED_BATCH_SIZE\"]]\n    seqs_batch = [test_seqs[pid] for pid in batch_ids]\n    # compute embeddings for this mini-batch\n    emb_mini = embed_batch_return_np(seqs_batch)  # shape (Bmini, D)\n    # append to buffer\n    embed_batch.append(emb_mini)\n    embed_ids.extend(batch_ids)\n    # if enough buffered to predict, or we're at the end, flush to prediction\n    buffered_examples = sum(arr.shape[0] for arr in embed_batch)\n    if buffered_examples >= CONFIG[\"PREDICT_BATCH_SIZE\"] or (i+CONFIG[\"EMBED_BATCH_SIZE\"] >= N_test):\n        # stack buffered embeddings (should be moderate size)\n        X_buffer = np.vstack(embed_batch).astype(np.float32)  # shape (Bbuf, D)\n        # predict in one shot for this buffer\n        y_buffer_prob = model.predict(X_buffer, batch_size=min(128, X_buffer.shape[0]), verbose=0)\n        # propagate per-batch (if desired)\n        if CONFIG[\"PROPAGATE_PREDICTIONS\"] and parents_map:\n            y_buffer_prob = propagate_batch(y_buffer_prob, restricted_parents, list(mlb.classes_), iterations=3)\n        # for each row, write top-K lines\n        for ridx, pid in enumerate(embed_ids):\n            probs = y_buffer_prob[ridx]\n            # pick top-K_PER_PROTEIN indices (and filter near-zero)\n            top_k = CONFIG[\"TOP_K_PER_PROTEIN\"]\n            if top_k is None:\n                idxs = np.where(probs >= best_thresh)[0]\n            else:\n                idxs = np.argsort(probs)[-top_k:]\n                idxs = [int(x) for x in idxs if probs[x] > 1e-6]\n            idxs = sorted(idxs, key=lambda x: probs[x], reverse=True)\n            for idx in idxs:\n                score = float(probs[idx])\n                if score <= 0.0: continue\n                go_id = mlb.classes_[idx]\n                out_f.write(f\"{pid}\\t{go_id}\\t{score:.3f}\\n\")\n        out_f.flush()\n        # free buffer\n        del X_buffer, y_buffer_prob, embed_batch\n        embed_batch = []\n        embed_ids = []\n        gc.collect()\n        if TORCH_AVAILABLE and torch.cuda.is_available(): torch.cuda.empty_cache()\n    # small progress print\n    if (i // CONFIG[\"EMBED_BATCH_SIZE\"]) % 50 == 0:\n        print(f\"[stream] processed {i} / {N_test}\")\n\nout_f.close()\nprint(f\"[done] Submission written to {CONFIG['OUTPUT_SUBMISSION']}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-08T23:11:21.423868Z","iopub.execute_input":"2025-12-08T23:11:21.424381Z","iopub.status.idle":"2025-12-08T23:15:10.789838Z","shell.execute_reply.started":"2025-12-08T23:11:21.424355Z","shell.execute_reply":"2025-12-08T23:15:10.788648Z"}},"outputs":[],"execution_count":null}]}