{"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":6180004,"sourceType":"datasetVersion","datasetId":3546496},{"sourceId":517313,"sourceType":"modelInstanceVersion","modelInstanceId":407954,"modelId":425816}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"forked from: \n- https://www.kaggle.com/code/subarnasaikia/cafa-6-base-model-lb-0-209/\n- https://www.kaggle.com/code/nihilisticneuralnet/protbert-ensemble/","metadata":{}},{"cell_type":"markdown","source":"# submission1","metadata":{}},{"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/submission1.tsv\"\n\nprint(\"Files are listed!!!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:53:49.752666Z","iopub.execute_input":"2025-10-16T22:53:49.75287Z","iopub.status.idle":"2025-10-16T22:53:49.761403Z","shell.execute_reply.started":"2025-10-16T22:53:49.752851Z","shell.execute_reply":"2025-10-16T22:53:49.760469Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Config","metadata":{}},{"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-10-16T22:53:49.767011Z","iopub.execute_input":"2025-10-16T22:53:49.767221Z","iopub.status.idle":"2025-10-16T22:53:49.777913Z","shell.execute_reply.started":"2025-10-16T22:53:49.767204Z","shell.execute_reply":"2025-10-16T22:53:49.777203Z"}},"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-10-16T22:53:49.779599Z","iopub.execute_input":"2025-10-16T22:53:49.779797Z","iopub.status.idle":"2025-10-16T22:53:57.927842Z","shell.execute_reply.started":"2025-10-16T22:53:49.779781Z","shell.execute_reply":"2025-10-16T22:53:57.927119Z"}},"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:53:57.928832Z","iopub.execute_input":"2025-10-16T22:53:57.929355Z","iopub.status.idle":"2025-10-16T22:53:57.940517Z","shell.execute_reply.started":"2025-10-16T22:53:57.929332Z","shell.execute_reply":"2025-10-16T22:53:57.939776Z"}},"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,"execution":{"iopub.status.busy":"2025-10-16T22:53:57.941163Z","iopub.execute_input":"2025-10-16T22:53:57.941369Z","iopub.status.idle":"2025-10-16T22:54:28.599019Z","shell.execute_reply.started":"2025-10-16T22:53:57.941352Z","shell.execute_reply":"2025-10-16T22:54:28.598343Z"}},"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_t33_650M_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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:54:28.600475Z","iopub.execute_input":"2025-10-16T22:54:28.600751Z","iopub.status.idle":"2025-10-16T22:54:28.654653Z","shell.execute_reply.started":"2025-10-16T22:54:28.600732Z","shell.execute_reply":"2025-10-16T22:54:28.654034Z"}},"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":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:54:28.655473Z","iopub.execute_input":"2025-10-16T22:54:28.6558Z","iopub.status.idle":"2025-10-16T22:54:42.580901Z","shell.execute_reply.started":"2025-10-16T22:54:28.655772Z","shell.execute_reply":"2025-10-16T22:54:42.579563Z"}},"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)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:54:42.581289Z","iopub.status.idle":"2025-10-16T22:54:42.581524Z","shell.execute_reply.started":"2025-10-16T22:54:42.581417Z","shell.execute_reply":"2025-10-16T22:54:42.581428Z"}},"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-10-16T22:54:42.582654Z","iopub.status.idle":"2025-10-16T22:54:42.582916Z","shell.execute_reply.started":"2025-10-16T22:54:42.582787Z","shell.execute_reply":"2025-10-16T22:54:42.582801Z"}},"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-10-16T22:54:42.583991Z","iopub.status.idle":"2025-10-16T22:54:42.584248Z","shell.execute_reply.started":"2025-10-16T22:54:42.584125Z","shell.execute_reply":"2025-10-16T22:54:42.584139Z"}},"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}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:54:42.585388Z","iopub.status.idle":"2025-10-16T22:54:42.585692Z","shell.execute_reply.started":"2025-10-16T22:54:42.585536Z","shell.execute_reply":"2025-10-16T22:54:42.58555Z"}},"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-10-16T22:54:42.586966Z","iopub.status.idle":"2025-10-16T22:54:42.587335Z","shell.execute_reply.started":"2025-10-16T22:54:42.587136Z","shell.execute_reply":"2025-10-16T22:54:42.587152Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# # ------------------------------------------------------------\n# # Save model / artifacts if desired (lightweight)\n# # ------------------------------------------------------------\n# model.save(\"/kaggle/working/cafa6_baseline_model\")\n# np.save(\"/kaggle/working/mlb_classes.npy\", np.array(mlb.classes_, dtype=object))\n# print(\"[done] saved model and classes; notebook finished.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-16T22:54:42.588538Z","iopub.status.idle":"2025-10-16T22:54:42.588761Z","shell.execute_reply.started":"2025-10-16T22:54:42.588658Z","shell.execute_reply":"2025-10-16T22:54:42.588667Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission1 = pd.read_csv(\"/kaggle/working/submission1.tsv\",\n                          sep='\\t', header=None, names=['Id', 'GO term', 'Confidence'])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# submission 2","metadata":{}},{"cell_type":"code","source":"submission2 = pd.read_csv('/kaggle/input/cafa5-055923-pred/submission.tsv',\n                          sep='\\t', header=None, names=['Id', 'GO term', 'Confidence']) ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# merge","metadata":{}},{"cell_type":"code","source":"submission1['Confidence'] = pd.to_numeric(submission1['Confidence'], errors='coerce')\nsubmission2['Confidence'] = pd.to_numeric(submission2['Confidence'], errors='coerce')\n\nmerged = submission1.merge(submission2,\n                           on=['Id', 'GO term'],\n                           how='outer',\n                           suffixes=('1', '2'))\n\nmerged['Confidence'] = merged['Confidence2'].combine_first(merged['Confidence1'])\nfinal_submission = merged[['Id', 'GO term', 'Confidence']]\nfinal_submission.to_csv('submission.tsv', sep='\\t', header=False, index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}