{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14066577,"sourceType":"datasetVersion","datasetId":8953533}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip -q install \"protobuf==5.29.4\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:06.369492Z","iopub.execute_input":"2025-12-13T09:12:06.369805Z","iopub.status.idle":"2025-12-13T09:12:09.625817Z","shell.execute_reply.started":"2025-12-13T09:12:06.369784Z","shell.execute_reply":"2025-12-13T09:12:09.625012Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, gc, random\nfrom collections import defaultdict, Counter\nfrom typing import Dict, List, Set, Tuple\n\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\nfrom scipy import sparse\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset\n\nfrom transformers import Trainer, TrainingArguments, EarlyStoppingCallback","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.627880Z","iopub.execute_input":"2025-12-13T09:12:09.628227Z","iopub.status.idle":"2025-12-13T09:12:09.634370Z","shell.execute_reply.started":"2025-12-13T09:12:09.628192Z","shell.execute_reply":"2025-12-13T09:12:09.633595Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SAMPLE_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\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\nGO_BASIC_OBO = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nOUTPUT_TSV = \"/kaggle/working/submission.tsv\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.635314Z","iopub.execute_input":"2025-12-13T09:12:09.635646Z","iopub.status.idle":"2025-12-13T09:12:09.654557Z","shell.execute_reply.started":"2025-12-13T09:12:09.635622Z","shell.execute_reply":"2025-12-13T09:12:09.653813Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CONFIG = {\n    \"TRAIN_FASTA\": TRAIN_SEQUENCES_FASTA,\n    \"TRAIN_TERMS\": TRAIN_TERMS_TSV,\n    \"GO_OBO\": GO_BASIC_OBO,\n    \"IA_FILE\": IA_TSV,\n    \"TEST_FASTA\": TESTSUPERSET_FASTA,\n    \"OUTPUT_SUBMISSION\": OUTPUT_TSV,\n\n    \"PROTBERT_TRAIN_EMB\": \"/kaggle/input/nnn-cafa6-protbert-embedding/train_embeddings.npy\",\n    \"PROTBERT_TEST_EMB\":  \"/kaggle/input/nnn-cafa6-protbert-embedding/test_embeddings.npy\",\n\n    \"TOP_K_LABELS\": 5000, \n\n    \"TFIDF_K\": 3, \n    \"TFIDF_NGRAM_RANGE\": (3, 4), # NEW \n    \"TFIDF_MAX_FEATURES\": 20000, \n    \"TFIDF_MIN_DF\": 2, \n    \"TFIDF_SUBLINEAR_TF\": True, # NEW\n\n    \"RANDOM_SEED\": 42,\n    \"TRAIN_BS\": 64,\n    \"EVAL_BS\": 128,\n    \"EPOCHS_TORCH\": 20,\n    \"LR_TORCH\": 5e-4,\n    \"PATIENCE\": 3,\n    \"WEIGHT_DECAY\": 1e-2,\n\n    \"PROJ_DIM\": 512,      \n    \"FUSION_HID\": 512,\n    \"DROPOUT_FUSION\": 0.3,\n    \"FP16\": True,\n\n    \"THRESH_COARSE\": [i/100 for i in range(1, 31)],\n    \"THRESH_REFINE_STEP\": 0.002,\n    \"THRESH_REFINE_SPAN\": 0.03,\n\n    \"TOP_K_PER_PROTEIN\": 200,\n\n    \"PROPAGATE_TRAIN_LABELS\": True,\n    \"PROPAGATE_PREDICTIONS\": True,\n    \"PROPAGATE_ITERATIONS\": 3,\n\n    \"EMBED_BATCH_SIZE\": 1024,\n    \"PREDICT_BATCH_SIZE\": 2048,  \n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.655988Z","iopub.execute_input":"2025-12-13T09:12:09.656177Z","iopub.status.idle":"2025-12-13T09:12:09.669167Z","shell.execute_reply.started":"2025-12-13T09:12:09.656163Z","shell.execute_reply":"2025-12-13T09:12:09.668418Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def set_seed(seed: int = 42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\nset_seed(CONFIG[\"RANDOM_SEED\"])\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"[env] device:\", device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.670075Z","iopub.execute_input":"2025-12-13T09:12:09.670289Z","iopub.status.idle":"2025-12-13T09:12:09.688499Z","shell.execute_reply.started":"2025-12-13T09:12:09.670267Z","shell.execute_reply":"2025-12-13T09:12:09.687769Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.689199Z","iopub.execute_input":"2025-12-13T09:12:09.690117Z","iopub.status.idle":"2025-12-13T09:12:09.703416Z","shell.execute_reply.started":"2025-12-13T09:12:09.690098Z","shell.execute_reply":"2025-12-13T09:12:09.702792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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():\n        mapping[r.protein].append(r.go)\n    print(f\"[io] Read training annotations for {len(mapping)} proteins from {path}\")\n    return mapping","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.704172Z","iopub.execute_input":"2025-12-13T09:12:09.704516Z","iopub.status.idle":"2025-12-13T09:12:09.722281Z","shell.execute_reply.started":"2025-12-13T09:12:09.704492Z","shell.execute_reply":"2025-12-13T09:12:09.721560Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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): \n        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]\": \n                cur_id=None\n            elif line.startswith(\"id: \"): \n                cur_id=line.split(\"id: \")[1].strip()\n            elif line.startswith(\"is_a: \"):\n                pid=line.split()[1].strip()\n                if cur_id: \n                    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: \n                        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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.723224Z","iopub.execute_input":"2025-12-13T09:12:09.723475Z","iopub.status.idle":"2025-12-13T09:12:09.739290Z","shell.execute_reply.started":"2025-12-13T09:12:09.723452Z","shell.execute_reply":"2025-12-13T09:12:09.738753Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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-12-13T09:12:09.739967Z","iopub.execute_input":"2025-12-13T09:12:09.740140Z","iopub.status.idle":"2025-12-13T09:12:09.757949Z","shell.execute_reply.started":"2025-12-13T09:12:09.740126Z","shell.execute_reply":"2025-12-13T09:12:09.757316Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from scipy import sparse as sp  # để dùng sp.issparse\n\nclass Evaluator:\n    def __init__(self, ia_weights: Dict[str, float], go_terms: List[str]):\n        self.go_terms = go_terms\n        self.weights = np.array([ia_weights.get(t, 1.0) for t in go_terms], dtype=np.float32)\n\n    def _dense_bool(self, y_true):\n        if sp.issparse(y_true):\n            return y_true.astype(np.bool_).toarray()\n        return (y_true > 0.5)\n\n    def _f_from_prec_rec(self, p: float, r: float, eps: float = 1e-12) -> float:\n        return (2.0 * p * r) / (p + r + eps)\n\n    def fmax(self, y_true, y_prob: np.ndarray, thresholds: np.ndarray) -> Tuple[float, float]:\n        y_true_b = self._dense_bool(y_true)\n        best_t, best_f = 0.5, -1.0\n\n        for t in thresholds:\n            y_pred_b = (y_prob >= t)\n            tp = np.logical_and(y_true_b, y_pred_b).sum(axis=1).astype(np.float32)\n            fp = np.logical_and(~y_true_b, y_pred_b).sum(axis=1).astype(np.float32)\n            fn = np.logical_and(y_true_b, ~y_pred_b).sum(axis=1).astype(np.float32)\n\n            has_pred = (tp + fp) > 0\n            has_true = (tp + fn) > 0\n\n            p = (tp[has_pred] / (tp[has_pred] + fp[has_pred] + 1e-12)).mean() if has_pred.any() else 0.0\n            r = (tp[has_true] / (tp[has_true] + fn[has_true] + 1e-12)).mean() if has_true.any() else 0.0\n\n            f = self._f_from_prec_rec(float(p), float(r))\n            if f > best_f:\n                best_f, best_t = f, float(t)\n\n        return best_t, best_f\n\n    def ia_fmax(self, y_true, y_prob: np.ndarray, thresholds: np.ndarray) -> Tuple[float, float]:\n        y_true_b = self._dense_bool(y_true)\n        w = self.weights[None, :]\n\n        best_t, best_f = 0.5, -1.0\n        for t in thresholds:\n            y_pred_b = (y_prob >= t)\n\n            tp_w = (np.logical_and(y_true_b, y_pred_b) * w).sum(axis=1).astype(np.float32)\n            fp_w = (np.logical_and(~y_true_b, y_pred_b) * w).sum(axis=1).astype(np.float32)\n            fn_w = (np.logical_and(y_true_b, ~y_pred_b) * w).sum(axis=1).astype(np.float32)\n\n            has_pred = (tp_w + fp_w) > 0\n            has_true = (tp_w + fn_w) > 0\n\n            p = (tp_w[has_pred] / (tp_w[has_pred] + fp_w[has_pred] + 1e-12)).mean() if has_pred.any() else 0.0\n            r = (tp_w[has_true] / (tp_w[has_true] + fn_w[has_true] + 1e-12)).mean() if has_true.any() else 0.0\n\n            f = self._f_from_prec_rec(float(p), float(r))\n            if f > best_f:\n                best_f, best_t = f, float(t)\n\n        return best_t, best_f\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.760268Z","iopub.execute_input":"2025-12-13T09:12:09.760612Z","iopub.status.idle":"2025-12-13T09:12:09.774118Z","shell.execute_reply.started":"2025-12-13T09:12:09.760594Z","shell.execute_reply":"2025-12-13T09:12:09.773478Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:09.774865Z","iopub.execute_input":"2025-12-13T09:12:09.775147Z","iopub.status.idle":"2025-12-13T09:12:11.567433Z","shell.execute_reply.started":"2025-12-13T09:12:09.775131Z","shell.execute_reply":"2025-12-13T09:12:11.566816Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_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\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\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):\n            extra |= get_ancestors(t, parents_map)\n        propagated[p]=sorted(terms|extra)\n    train_terms = propagated\n\nall_term_counts = Counter()\nfor p in train_proteins:\n    all_term_counts.update(train_terms[p])\n\nall_terms_sorted = [t for t,_ in all_term_counts.most_common()]\nchosen_terms = set(all_terms_sorted[:CONFIG[\"TOP_K_LABELS\"]])\nprint(f\"[prep] Restricting to top-{CONFIG['TOP_K_LABELS']} GO terms\")\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), sparse_output=True)\nY_csr = mlb.fit_transform(y_labels).astype(np.float32)  # sparse to save RAM\nprint(\"[prep] Label matrix shape (sparse):\", Y_csr.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:11.568303Z","iopub.execute_input":"2025-12-13T09:12:11.568561Z","iopub.status.idle":"2025-12-13T09:12:40.533908Z","shell.execute_reply.started":"2025-12-13T09:12:11.568540Z","shell.execute_reply":"2025-12-13T09:12:40.533182Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Embeddings: ProtBERT + TF-IDF","metadata":{}},{"cell_type":"code","source":"print(\"[emb] Loading ProtBERT embeddings...\")\nprot_train = np.load(CONFIG[\"PROTBERT_TRAIN_EMB\"]).astype(np.float32)\nprot_test  = np.load(CONFIG[\"PROTBERT_TEST_EMB\"]).astype(np.float32)\n\n# Assume aligned; slice to match train_proteins length\nprot_train = prot_train[:len(X_proteins)]\nprint(\"[emb] ProtBERT train:\", prot_train.shape, prot_train.dtype)\nprint(\"[emb] ProtBERT test :\", prot_test.shape, prot_test.dtype)\n\nprint(\"[emb] Building TF-IDF ...\")\ntfidf = TfidfVectorizer(\n    analyzer=\"char\",\n    ngram_range=CONFIG.get(\"TFIDF_NGRAM_RANGE\", (CONFIG[\"TFIDF_K\"], CONFIG[\"TFIDF_K\"])),\n    lowercase=False,\n    max_features=CONFIG[\"TFIDF_MAX_FEATURES\"],\n    min_df=CONFIG[\"TFIDF_MIN_DF\"],\n    sublinear_tf=CONFIG.get(\"TFIDF_SUBLINEAR_TF\", True),\n    dtype=np.float32,\n    norm=\"l2\",\n)\ntrain_seqs_list = [train_seqs[pid] for pid in X_proteins]\nX_tfidf_train = tfidf.fit_transform(train_seqs_list).astype(np.float32)\nprint(\"[emb] TF-IDF train:\", X_tfidf_train.shape, X_tfidf_train.dtype)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:12:40.534631Z","iopub.execute_input":"2025-12-13T09:12:40.534896Z","iopub.status.idle":"2025-12-13T09:13:34.010231Z","shell.execute_reply.started":"2025-12-13T09:12:40.534876Z","shell.execute_reply":"2025-12-13T09:13:34.009464Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train/val split by indices","metadata":{}},{"cell_type":"code","source":"idx_all = np.arange(len(X_proteins))\nidx_tr, idx_val = train_test_split(idx_all, test_size=0.15, random_state=CONFIG[\"RANDOM_SEED\"])\n\n# Dense y_val for metric\ny_val = Y_csr[idx_val].toarray().astype(np.float32)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:13:34.010958Z","iopub.execute_input":"2025-12-13T09:13:34.011209Z","iopub.status.idle":"2025-12-13T09:13:34.409983Z","shell.execute_reply.started":"2025-12-13T09:13:34.011188Z","shell.execute_reply":"2025-12-13T09:13:34.409388Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset + collate","metadata":{}},{"cell_type":"code","source":"class IndexDataset(Dataset):\n    def __init__(self, indices: np.ndarray):\n        self.indices = np.asarray(indices, dtype=np.int64)\n    def __len__(self): return len(self.indices)\n    def __getitem__(self, i):\n        return {\"idx\": int(self.indices[i])}\n\ndef fusion_collator(features):\n    idx = np.array([f[\"idx\"] for f in features], dtype=np.int64)\n\n    prot = torch.from_numpy(prot_train[idx]).to(torch.float32)  # (B, d_prot)\n\n    tfv = X_tfidf_train[idx].toarray().astype(np.float32)       # densify theo batch\n    tfv = torch.from_numpy(tfv)\n\n    labels = Y_csr[idx].toarray().astype(np.float32)\n    labels = torch.from_numpy(labels)\n\n    return {\"prot\": prot, \"tfidf\": tfv, \"labels\": labels}\n\ntrain_ds = IndexDataset(idx_tr)\nval_ds   = IndexDataset(idx_val)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:13:34.410651Z","iopub.execute_input":"2025-12-13T09:13:34.410881Z","iopub.status.idle":"2025-12-13T09:13:34.416757Z","shell.execute_reply.started":"2025-12-13T09:13:34.410847Z","shell.execute_reply":"2025-12-13T09:13:34.416109Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model: 2-tower fusion MLP","metadata":{}},{"cell_type":"code","source":"class FusionMLP(nn.Module):\n    def __init__(self, d_prot, d_tfidf, n_labels, proj_dim=512, hid=512, dropout=0.25):\n        super().__init__()\n        self.prot_proj = nn.Sequential(\n            nn.Linear(d_prot, proj_dim),\n            nn.BatchNorm1d(proj_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n        )\n        self.tfidf_proj = nn.Sequential(\n            nn.Linear(d_tfidf, proj_dim),\n            nn.BatchNorm1d(proj_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n        )\n        self.fuse = nn.Sequential(\n            nn.Linear(2 * proj_dim, hid),\n            nn.BatchNorm1d(hid),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hid, n_labels),  # logits\n        )\n\n    def forward(self, prot=None, tfidf=None, labels=None):\n        z1 = self.prot_proj(prot)\n        z2 = self.tfidf_proj(tfidf)\n        logits = self.fuse(torch.cat([z1, z2], dim=1))\n        return {\"logits\": logits}\n\nmodel = FusionMLP(\n    d_prot=prot_train.shape[1],\n    d_tfidf=X_tfidf_train.shape[1],\n    n_labels=Y_csr.shape[1],\n    proj_dim=CONFIG[\"PROJ_DIM\"],\n    hid=CONFIG[\"FUSION_HID\"],\n    dropout=CONFIG[\"DROPOUT_FUSION\"],\n).to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:13:34.417692Z","iopub.execute_input":"2025-12-13T09:13:34.417950Z","iopub.status.idle":"2025-12-13T09:13:34.527728Z","shell.execute_reply.started":"2025-12-13T09:13:34.417928Z","shell.execute_reply":"2025-12-13T09:13:34.527223Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train","metadata":{}},{"cell_type":"code","source":"class MultilabelTrainer(Trainer):\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.model_accepts_loss_kwargs = False  # :contentReference[oaicite:1]{index=1}\n\n    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None, **kwargs):\n        labels = inputs.pop(\"labels\")\n        outputs = model(**inputs)\n        logits = outputs[\"logits\"] if isinstance(outputs, dict) else outputs.logits\n        loss = F.binary_cross_entropy_with_logits(logits, labels)\n        return (loss, outputs) if return_outputs else loss\n\nargs = TrainingArguments(\n    output_dir=\"/kaggle/working/fusion_ckpt\",\n    do_train=True,\n    do_eval=True,\n\n    per_device_train_batch_size=CONFIG[\"TRAIN_BS\"],\n    per_device_eval_batch_size=CONFIG[\"EVAL_BS\"],\n    num_train_epochs=CONFIG[\"EPOCHS_TORCH\"],\n    learning_rate=CONFIG[\"LR_TORCH\"],\n    weight_decay=CONFIG[\"WEIGHT_DECAY\"],\n\n    fp16=CONFIG[\"FP16\"],\n    logging_steps=50,\n\n    eval_strategy=\"epoch\",\n    save_strategy=\"epoch\",\n    load_best_model_at_end=True,\n    metric_for_best_model=\"eval_loss\",\n    greater_is_better=False,\n\n    remove_unused_columns=False,\n    report_to=\"none\",\n    lr_scheduler_type=\"cosine\",\n    warmup_ratio=0.05,\n    max_grad_norm=1.0,\n)\n\ntrainer = MultilabelTrainer(\n    model=model,\n    args=args,\n    train_dataset=train_ds,\n    eval_dataset=val_ds,\n    data_collator=fusion_collator,\n    callbacks=[EarlyStoppingCallback(early_stopping_patience=CONFIG[\"PATIENCE\"])],\n)\n\ntrainer.train()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T09:13:34.528382Z","iopub.execute_input":"2025-12-13T09:13:34.528540Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pred = trainer.predict(val_ds)\nval_logits = pred.predictions  # numpy\ny_val_prob = 1.0 / (1.0 + np.exp(-val_logits))  # sigmoid","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"term_to_idx = {t:i for i,t in enumerate(mlb.classes_)}\nclasses_list = list(mlb.classes_)\n\nrestricted_parents = {}\nfor t in classes_list:\n    restricted_parents[t] = set([p for p in parents_map.get(t, set()) if p in term_to_idx])\n\nparent_idx_list = []\nfor child_term in classes_list:\n    parent_idx_list.append([term_to_idx[p] for p in restricted_parents.get(child_term, [])])\n\ndef propagate_batch(pred_batch: np.ndarray, parent_idx_list, iterations=3):\n    for _ in range(iterations):\n        changed = False\n        for child_idx in range(pred_batch.shape[1]):\n            child_scores = pred_batch[:, child_idx]\n            for pidx in parent_idx_list[child_idx]:\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:\n            break\n    return pred_batch","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"y_true_val = Y_csr[idx_val]  # (N_val, L) sparse\n\ny_eval_prob = y_val_prob\nif CONFIG[\"PROPAGATE_PREDICTIONS\"] and parents_map:\n    y_eval_prob = propagate_batch(\n        y_eval_prob.copy(),\n        parent_idx_list,\n        iterations=CONFIG[\"PROPAGATE_ITERATIONS\"]\n    )\n\nevaluator = Evaluator(ia_weights, list(mlb.classes_))\n\nthr_coarse = np.array(CONFIG[\"THRESH_COARSE\"], dtype=np.float32)\nbest_thresh, best_score = evaluator.ia_fmax(y_true_val, y_eval_prob, thr_coarse)\n\nlo = max(0.0, best_thresh - CONFIG[\"THRESH_REFINE_SPAN\"])\nhi = min(0.999, best_thresh + CONFIG[\"THRESH_REFINE_SPAN\"])\nthr_refine = np.arange(lo, hi + 1e-12, CONFIG[\"THRESH_REFINE_STEP\"], dtype=np.float32)\nbest_thresh, best_score = evaluator.ia_fmax(y_true_val, y_eval_prob, thr_refine)\n\nprint(f\"[eval] best_thresh {best_thresh:.3f} best IA-weighted Fmax {best_score:.6f}\")\n\nt_f, f_f = evaluator.fmax(y_true_val, y_eval_prob, thr_refine)\nprint(f\"[eval] best_thresh(Fmax) {t_f:.3f} best Fmax {f_f:.6f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"out_fpath = CONFIG[\"OUTPUT_SUBMISSION\"]\nopen(out_fpath, \"w\").close()\nout_f = open(out_fpath, \"a\")\n\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']} / predict {CONFIG['PREDICT_BATCH_SIZE']}\")\n\ndef tfidf_transform_batch(seq_list: List[str]) -> np.ndarray:\n    return tfidf.transform(seq_list).toarray().astype(np.float32)\n\nembed_plm_buf = []\nembed_tfidf_buf = []\nembed_ids_buf = []\n\ntrainer.model.to(device)\ntrainer.model.eval()\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\n    # ProtBERT slice\n    plm_mini = prot_test[i:i+len(batch_ids)]\n    tfidf_mini = tfidf_transform_batch(seqs_batch)\n\n    embed_plm_buf.append(plm_mini)\n    embed_tfidf_buf.append(tfidf_mini)\n    embed_ids_buf.extend(batch_ids)\n\n    buffered = sum(a.shape[0] for a in embed_plm_buf)\n    if buffered >= CONFIG[\"PREDICT_BATCH_SIZE\"] or (i+CONFIG[\"EMBED_BATCH_SIZE\"] >= N_test):\n        Xp = np.vstack(embed_plm_buf).astype(np.float32)\n        Xt = np.vstack(embed_tfidf_buf).astype(np.float32)\n\n        with torch.no_grad(), torch.cuda.amp.autocast(enabled=(device.type==\"cuda\")):\n            xb1 = torch.from_numpy(Xp).to(device)\n            xb2 = torch.from_numpy(Xt).to(device)\n            logits = trainer.model(prot=xb1, tfidf=xb2)[\"logits\"]\n            probs = torch.sigmoid(logits).float().cpu().numpy()\n\n        if CONFIG[\"PROPAGATE_PREDICTIONS\"] and parents_map:\n            probs = propagate_batch(probs, parent_idx_list, iterations=CONFIG[\"PROPAGATE_ITERATIONS\"])\n\n        top_k = CONFIG[\"TOP_K_PER_PROTEIN\"]\n        for ridx, pid in enumerate(embed_ids_buf):\n            row = probs[ridx]\n            idxs = np.argsort(row)[-top_k:]\n            idxs = [int(x) for x in idxs if row[x] > 1e-6]\n            idxs = sorted(idxs, key=lambda x: row[x], reverse=True)\n            for idx in idxs:\n                score = float(row[idx])\n                if score <= 0.0: \n                    continue\n                go_id = mlb.classes_[idx]\n                out_f.write(f\"{pid}\\t{go_id}\\t{score:.6f}\\n\")\n\n        out_f.flush()\n        embed_plm_buf, embed_tfidf_buf, embed_ids_buf = [], [], []\n        del Xp, Xt, probs, logits, xb1, xb2\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n\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},"outputs":[],"execution_count":null}]}