{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14066577,"sourceType":"datasetVersion","datasetId":8953533},{"sourceId":14066843,"sourceType":"datasetVersion","datasetId":8950709},{"sourceId":14077114,"sourceType":"datasetVersion","datasetId":8961200},{"sourceId":14102744,"sourceType":"datasetVersion","datasetId":8981800}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install protobuf==4.25.3","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:38:51.994087Z","iopub.execute_input":"2025-12-09T18:38:51.994285Z","iopub.status.idle":"2025-12-09T18:38:57.414706Z","shell.execute_reply.started":"2025-12-09T18:38:51.994268Z","shell.execute_reply":"2025-12-09T18:38:57.413977Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport gc\nimport random\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nfrom typing import Optional\nfrom typing import Dict, List, Set, Tuple\nfrom tqdm.auto import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.model_selection import train_test_split","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","execution":{"iopub.status.busy":"2025-12-09T18:38:57.416980Z","iopub.execute_input":"2025-12-09T18:38:57.417227Z","iopub.status.idle":"2025-12-09T18:39:02.175626Z","shell.execute_reply.started":"2025-12-09T18:38:57.417203Z","shell.execute_reply":"2025-12-09T18:39:02.174897Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# CONFIGURATION","metadata":{}},{"cell_type":"code","source":"# Paths\nDATA_DIR = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\nTRAIN_DIR = DATA_DIR / \"Train\"\nTEST_DIR = DATA_DIR / \"Test\"\nWORK_DIR = Path(\"/kaggle/working\")\n\nTRAIN_FASTA = TRAIN_DIR / \"train_sequences.fasta\"\nTRAIN_TERMS = TRAIN_DIR / \"train_terms.tsv\"\nGO_OBO = TRAIN_DIR / \"go-basic.obo\"\nIA_FILE = DATA_DIR / \"IA.tsv\"\nTEST_FASTA = TEST_DIR / \"testsuperset.fasta\"\nOUTPUT_FILE = WORK_DIR / \"submission.tsv\"\n\n# ProtBERT embeddings (cập nhật đường dẫn của bạn)\nPROTBERT_TRAIN_EMB = Path(\"/kaggle/input/nnn-cafa6-protbert-embedding/train_embeddings.npy\")\nPROTBERT_TEST_EMB = Path(\"/kaggle/input/nnn-cafa6-protbert-embedding/test_embeddings.npy\")\n\n# Test split\nTEST_PROPORTION = 0.15\n\n# Model parameters\nRANDOM_SEED = 42\nTOP_K_LABELS = 3000\nHIDDEN_DIMS = [1024, 512]\nDROPOUT = 0.4\nLEARNING_RATE = 1e-3\nBATCH_SIZE = 32\nEPOCHS = 25\nPATIENCE = 5\n\n# Prediction parameters\nTOP_K_PER_PROTEIN = 200\nTHRESHOLD_SEARCH = True\nTHRESHOLD_GRID = np.arange(0.01, 0.51, 0.01)\n\n# GO propagation\nPROPAGATE_TRAIN = True\nPROPAGATE_PRED = True\nPROPAGATE_ITERATIONS = 3\n\n# Device\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {DEVICE}\")\n\n# Set seeds\nrandom.seed(RANDOM_SEED)\nnp.random.seed(RANDOM_SEED)\ntorch.manual_seed(RANDOM_SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(RANDOM_SEED)","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.176469Z","iopub.execute_input":"2025-12-09T18:39:02.176868Z","iopub.status.idle":"2025-12-09T18:39:02.263538Z","shell.execute_reply.started":"2025-12-09T18:39:02.176839Z","shell.execute_reply":"2025-12-09T18:39:02.262901Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data Helper Functions","metadata":{}},{"cell_type":"code","source":"def read_fasta(path: Path) -> Dict[str, str]:\n    \"\"\"Read FASTA file and return dict of protein_id: sequence\"\"\"\n    sequences = {}\n    with open(path) as f:\n        protein_id = None\n        seq_parts = []\n        \n        for line in f:\n            line = line.strip()\n            if line.startswith(\">\"):\n                if protein_id:\n                    sequences[protein_id] = \"\".join(seq_parts)\n                \n                header = line[1:].split()[0]\n                protein_id = header.split(\"|\")[1] if \"|\" in header else header\n                seq_parts = []\n            else:\n                seq_parts.append(line)\n        \n        if protein_id:\n            sequences[protein_id] = \"\".join(seq_parts)\n    \n    print(f\"✓ Loaded {len(sequences):,} sequences from {path.name}\")\n    return sequences\n\n\ndef read_annotations(path: Path) -> Dict[str, List[str]]:\n    \"\"\"Read protein-GO term annotations\"\"\"\n    df = pd.read_csv(path, sep=\"\\t\", header=None, \n                    names=[\"protein\", \"go_term\", \"ontology\"])\n    \n    annotations = defaultdict(list)\n    for _, row in df.iterrows():\n        annotations[row.protein].append(row.go_term)\n    \n    print(f\"✓ Loaded annotations for {len(annotations):,} proteins\")\n    return dict(annotations)\n\n\ndef read_ia_weights(path: Path) -> Dict[str, float]:\n    \"\"\"Read Information Accretion weights\"\"\"\n    if not path.exists():\n        print(\"⚠ IA weights file not found\")\n        return {}\n    \n    df = pd.read_csv(path, sep=\"\\t\", header=None, names=[\"go_term\", \"ia\"])\n    weights = {}\n    \n    for _, row in df.iterrows():\n        try:\n            weights[row.go_term] = float(str(row.ia).replace(\",\", \".\"))\n        except:\n            weights[row.go_term] = 0.0\n    \n    print(f\"✓ Loaded IA weights for {len(weights):,} GO terms\")\n    return weights","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.264190Z","iopub.execute_input":"2025-12-09T18:39:02.264477Z","iopub.status.idle":"2025-12-09T18:39:02.277076Z","shell.execute_reply.started":"2025-12-09T18:39:02.264451Z","shell.execute_reply":"2025-12-09T18:39:02.276385Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# GO ontology Helper Functions","metadata":{}},{"cell_type":"code","source":"def parse_obo(path: Path):\n    \"\"\"Parse OBO file to extract parent-child relationships + namespace (BP/MF/CC)\"\"\"\n    parents = defaultdict(set)\n    children = defaultdict(set)\n    go_namespace = {}   # go_id -> \"BP\"/\"MF\"/\"CC\"\n    \n    if not path.exists():\n        print(\"⚠ OBO file not found\")\n        return parents, children, go_namespace\n    \n    with open(path) as f:\n        current_id = None\n        \n        for line in f:\n            line = line.strip()\n            \n            if line == \"[Term]\":\n                current_id = None\n            elif line.startswith(\"id: \"):\n                current_id = line.split(\"id: \")[1]\n            elif line.startswith(\"namespace:\") and current_id:\n                ns = line.split(\"namespace:\")[1].strip()\n                if ns == \"biological_process\":\n                    go_namespace[current_id] = \"BP\"\n                elif ns == \"molecular_function\":\n                    go_namespace[current_id] = \"MF\"\n                elif ns == \"cellular_component\":\n                    go_namespace[current_id] = \"CC\"\n            elif line.startswith(\"is_a: \") and current_id:\n                parent_id = line.split()[1]\n                parents[current_id].add(parent_id)\n                children[parent_id].add(current_id)\n            elif line.startswith(\"relationship: part_of \") and current_id:\n                parts = line.split()\n                if len(parts) >= 3:\n                    parent_id = parts[2]\n                    parents[current_id].add(parent_id)\n                    children[parent_id].add(current_id)\n    \n    print(f\"✓ Parsed GO graph: {len(parents):,} terms with parents\")\n    print(f\"✓ Parsed namespaces for {len(go_namespace):,} GO terms\")\n    return dict(parents), dict(children), go_namespace\n\n\n\ndef get_ancestors(go_term: str, parents: Dict) -> Set[str]:\n    \"\"\"Get all ancestor terms\"\"\"\n    ancestors = set()\n    stack = [go_term]\n    \n    while stack:\n        current = stack.pop()\n        for parent in parents.get(current, []):\n            if parent not in ancestors:\n                ancestors.add(parent)\n                stack.append(parent)\n    \n    return ancestors\n\n\ndef propagate_labels(annotations: Dict[str, List[str]], \n                    parents: Dict) -> Dict[str, List[str]]:\n    \"\"\"Propagate labels up the GO graph\"\"\"\n    print(\"→ Propagating labels up GO hierarchy...\")\n    propagated = {}\n    \n    for protein, terms in tqdm(annotations.items(), desc=\"Propagating\"):\n        expanded = set(terms)\n        for term in terms:\n            expanded.update(get_ancestors(term, parents))\n        propagated[protein] = sorted(expanded)\n    \n    original_count = sum(len(v) for v in annotations.values())\n    new_count = sum(len(v) for v in propagated.values())\n    print(f\"  {original_count:,} → {new_count:,} annotations\")\n    \n    return propagated","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.277939Z","iopub.execute_input":"2025-12-09T18:39:02.278721Z","iopub.status.idle":"2025-12-09T18:39:02.296896Z","shell.execute_reply.started":"2025-12-09T18:39:02.278684Z","shell.execute_reply":"2025-12-09T18:39:02.296288Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Pytorch Dataset","metadata":{}},{"cell_type":"code","source":"class ProteinDataset(Dataset):\n    \"\"\"Dataset cho multi-task MF/BP/CC\"\"\"\n    def __init__(self, embeddings, y_bp, y_mf, y_cc):\n        self.embeddings = torch.FloatTensor(embeddings)\n        self.y_bp = torch.FloatTensor(y_bp)\n        self.y_mf = torch.FloatTensor(y_mf)\n        self.y_cc = torch.FloatTensor(y_cc)\n\n    def __len__(self):\n        return len(self.embeddings)\n\n    def __getitem__(self, idx):\n        return (\n            self.embeddings[idx],\n            self.y_bp[idx],\n            self.y_mf[idx],\n            self.y_cc[idx],\n        )","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.297603Z","iopub.execute_input":"2025-12-09T18:39:02.297868Z","iopub.status.idle":"2025-12-09T18:39:02.315035Z","shell.execute_reply.started":"2025-12-09T18:39:02.297843Z","shell.execute_reply":"2025-12-09T18:39:02.314354Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Pytorch Model","metadata":{}},{"cell_type":"code","source":"class MultiTaskProteinFunctionPredictor(nn.Module):\n    \"\"\"MLP multi-task cho BP / MF / CC\"\"\"\n\n    def __init__(\n        self,\n        input_dim: int,\n        bp_dim: int,\n        mf_dim: int,\n        cc_dim: int,\n        hidden_dims: List[int] = [512, 256],\n        dropout: float = 0.3,\n    ):\n        super().__init__()\n\n        layers = []\n        prev_dim = input_dim\n\n        # shared trunk\n        for hidden_dim in hidden_dims:\n            layers.append(nn.Linear(prev_dim, hidden_dim))\n            layers.append(nn.BatchNorm1d(hidden_dim))\n            layers.append(nn.ReLU())\n            layers.append(nn.Dropout(dropout))\n            prev_dim = hidden_dim\n\n        self.trunk = nn.Sequential(*layers)\n\n        # 3 head output\n        self.bp_head = nn.Linear(prev_dim, bp_dim)\n        self.mf_head = nn.Linear(prev_dim, mf_dim)\n        self.cc_head = nn.Linear(prev_dim, cc_dim)\n\n    def forward(self, x):\n        h = self.trunk(x)\n        out_bp = self.bp_head(h)\n        out_mf = self.mf_head(h)\n        out_cc = self.cc_head(h)\n        return out_bp, out_mf, out_cc","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.317474Z","iopub.execute_input":"2025-12-09T18:39:02.317738Z","iopub.status.idle":"2025-12-09T18:39:02.328467Z","shell.execute_reply.started":"2025-12-09T18:39:02.317722Z","shell.execute_reply":"2025-12-09T18:39:02.327800Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training Settings","metadata":{}},{"cell_type":"code","source":"def train_epoch(\n    model,\n    dataloader,\n    criterion_bp,\n    criterion_mf,\n    criterion_cc,\n    optimizer,\n    device,\n    w_bp=1.0,\n    w_mf=1.0,\n    w_cc=1.0,\n):\n    model.train()\n    total_loss = 0.0\n\n    for embeddings, y_bp, y_mf, y_cc in dataloader:\n        embeddings = embeddings.to(device)\n        y_bp = y_bp.to(device)\n        y_mf = y_mf.to(device)\n        y_cc = y_cc.to(device)\n\n        optimizer.zero_grad()\n        logits_bp, logits_mf, logits_cc = model(embeddings)\n\n        loss_bp = criterion_bp(logits_bp, y_bp)\n        loss_mf = criterion_mf(logits_mf, y_mf)\n        loss_cc = criterion_cc(logits_cc, y_cc)\n\n        loss = w_bp * loss_bp + w_mf * loss_mf + w_cc * loss_cc\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n\n    return total_loss / len(dataloader)\n\n\ndef evaluate(\n    model,\n    dataloader,\n    criterion_bp,\n    criterion_mf,\n    criterion_cc,\n    device,\n    w_bp=1.0,\n    w_mf=1.0,\n    w_cc=1.0,\n):\n    model.eval()\n    total_loss = 0.0\n\n    all_bp_preds, all_mf_preds, all_cc_preds = [], [], []\n    all_bp_labels, all_mf_labels, all_cc_labels = [], [], []\n\n    with torch.no_grad():\n        for embeddings, y_bp, y_mf, y_cc in dataloader:\n            embeddings = embeddings.to(device)\n            y_bp = y_bp.to(device)\n            y_mf = y_mf.to(device)\n            y_cc = y_cc.to(device)\n\n            logits_bp, logits_mf, logits_cc = model(embeddings)\n\n            loss_bp = criterion_bp(logits_bp, y_bp)\n            loss_mf = criterion_mf(logits_mf, y_mf)\n            loss_cc = criterion_cc(logits_cc, y_cc)\n\n            loss = w_bp * loss_bp + w_mf * loss_mf + w_cc * loss_cc\n            total_loss += loss.item()\n\n            prob_bp = torch.sigmoid(logits_bp)\n            prob_mf = torch.sigmoid(logits_mf)\n            prob_cc = torch.sigmoid(logits_cc)\n\n            all_bp_preds.append(prob_bp.cpu().numpy())\n            all_mf_preds.append(prob_mf.cpu().numpy())\n            all_cc_preds.append(prob_cc.cpu().numpy())\n\n            all_bp_labels.append(y_bp.cpu().numpy())\n            all_mf_labels.append(y_mf.cpu().numpy())\n            all_cc_labels.append(y_cc.cpu().numpy())\n\n    avg_loss = total_loss / len(dataloader)\n\n    preds_bp = np.vstack(all_bp_preds)\n    preds_mf = np.vstack(all_mf_preds)\n    preds_cc = np.vstack(all_cc_preds)\n\n    labels_bp = np.vstack(all_bp_labels)\n    labels_mf = np.vstack(all_mf_labels)\n    labels_cc = np.vstack(all_cc_labels)\n\n    return avg_loss, preds_bp, preds_mf, preds_cc, labels_bp, labels_mf, labels_cc\n\n\ndef compute_metrics(\n    y_true: np.ndarray,\n    y_pred: np.ndarray,\n    weights: Optional[np.ndarray] = None,\n) -> Dict:\n    \"\"\"Compute (có thể) IA-weighted precision, recall, F1.\n       Nếu weights=None → tính trung bình không IA.\n    \"\"\"\n    tp = ((y_true == 1) & (y_pred == 1)).sum(axis=0).astype(float)\n    fp = ((y_true == 0) & (y_pred == 1)).sum(axis=0).astype(float)\n    fn = ((y_true == 1) & (y_pred == 0)).sum(axis=0).astype(float)\n    \n    eps = 1e-12\n    precision = tp / (tp + fp + eps)\n    recall = tp / (tp + fn + eps)\n    f1 = 2 * precision * recall / (precision + recall + eps)\n    \n    if weights is None:\n        # Không IA: coi mọi class đóng góp như nhau\n        weights = np.ones_like(precision)\n    \n    weight_sum = weights.sum() + eps\n    return {\n        \"precision\": (precision * weights).sum() / weight_sum,\n        \"recall\": (recall * weights).sum() / weight_sum,\n        \"f1\": (f1 * weights).sum() / weight_sum,\n    }\n\n\ndef find_best_threshold(y_true: np.ndarray, y_pred_prob: np.ndarray,\n                       weights: np.ndarray, thresholds: np.ndarray) -> Tuple[float, float]:\n    \"\"\"Find optimal threshold by grid search\"\"\"\n    print(\"→ Searching for optimal threshold...\")\n    \n    best_threshold = 0.5\n    best_f1 = -1.0\n    \n    for thresh in tqdm(thresholds, desc=\"Threshold search\"):\n        y_pred = (y_pred_prob >= thresh).astype(int)\n        metrics = compute_metrics(y_true, y_pred, weights)\n        \n        if metrics[\"f1\"] > best_f1:\n            best_f1 = metrics[\"f1\"]\n            best_threshold = thresh\n\n    print(f\"  Best threshold: {best_threshold:.3f} (F1: {best_f1:.4f})\")\n    return best_threshold, best_f1","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:46:28.622816Z","iopub.execute_input":"2025-12-09T18:46:28.623503Z","iopub.status.idle":"2025-12-09T18:46:28.636352Z","shell.execute_reply.started":"2025-12-09T18:46:28.623481Z","shell.execute_reply":"2025-12-09T18:46:28.635711Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prediction Propagation","metadata":{}},{"cell_type":"code","source":"def propagate_predictions(predictions: np.ndarray, \n                         parents: Dict,\n                         go_terms: List[str],\n                         iterations: int = 3) -> np.ndarray:\n    \"\"\"Propagate predictions up GO hierarchy\"\"\"\n    print(\"→ Propagating predictions...\")\n    \n    term_to_idx = {t: i for i, t in enumerate(go_terms)}\n    pred_copy = predictions.copy()\n    \n    # Restrict parents to terms in vocabulary\n    restricted_parents = {}\n    for term in go_terms:\n        restricted_parents[term] = {p for p in parents.get(term, []) \n                                   if p in term_to_idx}\n    \n    for iteration in range(iterations):\n        changed = False\n        \n        for child_idx, child_term in enumerate(go_terms):\n            child_scores = pred_copy[:, child_idx]\n            \n            for parent_term in restricted_parents.get(child_term, []):\n                parent_idx = term_to_idx[parent_term]\n                \n                # Update parent where child score is higher\n                mask = child_scores > pred_copy[:, parent_idx]\n                if mask.any():\n                    pred_copy[mask, parent_idx] = child_scores[mask]\n                    changed = True\n        \n        if not changed:\n            print(f\"  Converged after {iteration + 1} iterations\")\n            break\n    \n    return pred_copy","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.349846Z","iopub.execute_input":"2025-12-09T18:39:02.350035Z","iopub.status.idle":"2025-12-09T18:39:02.365314Z","shell.execute_reply.started":"2025-12-09T18:39:02.350021Z","shell.execute_reply":"2025-12-09T18:39:02.364668Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main Pipeline","metadata":{}},{"cell_type":"code","source":"# 1. Load Data\nprint(\"📁 Loading data...\")\ntrain_seqs = read_fasta(TRAIN_FASTA)\ntest_seqs = read_fasta(TEST_FASTA)\nannotations = read_annotations(TRAIN_TERMS)\nia_weights_dict = read_ia_weights(IA_FILE)\n\nprint(\"\\n🌳 Loading GO ontology...\")\nparents_map, children_map, go_namespace = parse_obo(GO_OBO)","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:02.365926Z","iopub.execute_input":"2025-12-09T18:39:02.366111Z","iopub.status.idle":"2025-12-09T18:39:29.605546Z","shell.execute_reply.started":"2025-12-09T18:39:02.366097Z","shell.execute_reply":"2025-12-09T18:39:29.604928Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2. Propagate Training Labels\nif PROPAGATE_TRAIN:\n    annotations = propagate_labels(annotations, parents_map)","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:29.606279Z","iopub.execute_input":"2025-12-09T18:39:29.606501Z","iopub.status.idle":"2025-12-09T18:39:32.478433Z","shell.execute_reply.started":"2025-12-09T18:39:29.606478Z","shell.execute_reply":"2025-12-09T18:39:32.477632Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3. Prepare Labels\nprint(\"\\n🏷️  Preparing labels (multi-task BP/MF/CC)...\")\n\n# Get proteins with both sequences and annotations\ntrain_proteins = [p for p in annotations.keys() if p in train_seqs]\nprint(f\"  {len(train_proteins):,} training proteins\")\n\n# Count term frequencies (trong tất cả ontology)\nterm_counts = Counter()\nfor protein in train_proteins:\n    term_counts.update(annotations[protein])\n\n# Select top-K terms (giống như trước)\ntop_terms = [t for t, _ in term_counts.most_common(TOP_K_LABELS)]\nchosen_terms = set(top_terms)\nprint(f\"  Using top {len(chosen_terms):,} GO terms (all ontologies combined)\")\n\n# Filter annotations chỉ giữ các term trong top-K\nfor protein in train_proteins:\n    annotations[protein] = [t for t in annotations[protein] if t in chosen_terms]\n\n# Tách top_terms theo ontology dùng go_namespace từ parse_obo\nbp_terms = [t for t in top_terms if go_namespace.get(t) == \"BP\"]\nmf_terms = [t for t in top_terms if go_namespace.get(t) == \"MF\"]\ncc_terms = [t for t in top_terms if go_namespace.get(t) == \"CC\"]\n\nbp_term_set = set(bp_terms)\nmf_term_set = set(mf_terms)\ncc_term_set = set(cc_terms)\n\nprint(f\"  BP terms: {len(bp_terms)}\")\nprint(f\"  MF terms: {len(mf_terms)}\")\nprint(f\"  CC terms: {len(cc_terms)}\")\nprint(f\"  Total (BP+MF+CC): {len(bp_terms) + len(mf_terms) + len(cc_terms)}\")\n\n# Tạo list label cho từng ontology\nlabels_bp = [[t for t in annotations[p] if t in bp_term_set] for p in train_proteins]\nlabels_mf = [[t for t in annotations[p] if t in mf_term_set] for p in train_proteins]\nlabels_cc = [[t for t in annotations[p] if t in cc_term_set] for p in train_proteins]\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\n\nmlb_bp = MultiLabelBinarizer(classes=sorted(bp_term_set))\nmlb_mf = MultiLabelBinarizer(classes=sorted(mf_term_set))\nmlb_cc = MultiLabelBinarizer(classes=sorted(cc_term_set))\n\ny_bp = mlb_bp.fit_transform(labels_bp).astype(np.float32)\ny_mf = mlb_mf.fit_transform(labels_mf).astype(np.float32)\ny_cc = mlb_cc.fit_transform(labels_cc).astype(np.float32)\n\nprint(f\"  y_bp shape: {y_bp.shape}\")\nprint(f\"  y_mf shape: {y_mf.shape}\")\nprint(f\"  y_cc shape: {y_cc.shape}\")\n\n# IA weights cho từng ontology\nia_weights_bp = np.array([ia_weights_dict.get(t, 1.0) for t in mlb_bp.classes_])\nia_weights_mf = np.array([ia_weights_dict.get(t, 1.0) for t in mlb_mf.classes_])\nia_weights_cc = np.array([ia_weights_dict.get(t, 1.0) for t in mlb_cc.classes_])","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:32.479380Z","iopub.execute_input":"2025-12-09T18:39:32.479751Z","iopub.status.idle":"2025-12-09T18:39:36.387292Z","shell.execute_reply.started":"2025-12-09T18:39:32.479725Z","shell.execute_reply":"2025-12-09T18:39:36.386629Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === CÂN LẠI TRỌNG SỐ LOSS CHO 3 ONTOLOGY ===\nn_bp = y_bp.shape[1]\nn_mf = y_mf.shape[1]\nn_cc = y_cc.shape[1]\n\nimport math\ninv_sqrt = np.array([\n    1.0 / math.sqrt(n_bp),\n    1.0 / math.sqrt(n_mf),\n    1.0 / math.sqrt(n_cc),\n])\ninv_sqrt = inv_sqrt / inv_sqrt.sum()\nW_BP, W_MF, W_CC = inv_sqrt.tolist()\n\nprint(\"Loss weights:\",\n      f\"W_BP={W_BP:.3f}, W_MF={W_MF:.3f}, W_CC={W_CC:.3f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 4. Load ProtBERT Embeddings\nprint(\"\\n🧬 Loading ProtBERT embeddings...\")\n\n# Load train embeddings\ntrain_emb = np.load(PROTBERT_TRAIN_EMB).astype(np.float32)\nprint(f\"  Train embeddings shape: {train_emb.shape}\")\n\n# TODO: Align embeddings with protein IDs if needed\n# Giả sử embeddings đã được align với train_proteins\nX = train_emb[:len(train_proteins)]","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:36.388046Z","iopub.execute_input":"2025-12-09T18:39:36.388315Z","iopub.status.idle":"2025-12-09T18:39:38.556091Z","shell.execute_reply.started":"2025-12-09T18:39:36.388295Z","shell.execute_reply":"2025-12-09T18:39:38.555256Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 5. Train/Val Split\nprint(\"\\n📊 Splitting data (multi-task)...\")\nX_train, X_val, y_bp_train, y_bp_val, y_mf_train, y_mf_val, y_cc_train, y_cc_val = train_test_split(\n    X,\n    y_bp,\n    y_mf,\n    y_cc,\n    test_size=TEST_PROPORTION,\n    random_state=RANDOM_SEED,\n)\n\nprint(f\"  Train: {X_train.shape}\")\nprint(f\"  Val:   {X_val.shape}\")\n\ntrain_dataset = ProteinDataset(X_train, y_bp_train, y_mf_train, y_cc_train)\nval_dataset = ProteinDataset(X_val, y_bp_val, y_mf_val, y_cc_val)\n\ntrain_loader = DataLoader(\n    train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True\n)\nval_loader = DataLoader(\n    val_dataset, batch_size=BATCH_SIZE * 2, shuffle=False, num_workers=2, pin_memory=True\n)","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:39:38.556919Z","iopub.execute_input":"2025-12-09T18:39:38.557111Z","iopub.status.idle":"2025-12-09T18:39:38.991508Z","shell.execute_reply.started":"2025-12-09T18:39:38.557096Z","shell.execute_reply":"2025-12-09T18:39:38.990706Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n🏗️  Building multi-task model...\")\n\nmodel = MultiTaskProteinFunctionPredictor(\n    input_dim=X.shape[1],\n    bp_dim=y_bp.shape[1],\n    mf_dim=y_mf.shape[1],\n    cc_dim=y_cc.shape[1],\n    hidden_dims=HIDDEN_DIMS,\n    dropout=DROPOUT,\n).to(DEVICE)\n\ntotal_params = sum(p.numel() for p in model.parameters())\ntrainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\nprint(f\"  Total parameters: {total_params:,}\")\nprint(f\"  Trainable parameters: {trainable_params:,}\")\n\n# loss cho từng head\ncriterion_bp = nn.BCEWithLogitsLoss()\ncriterion_mf = nn.BCEWithLogitsLoss()\ncriterion_cc = nn.BCEWithLogitsLoss()\n\noptimizer = Adam(model.parameters(), lr=LEARNING_RATE)\nscheduler = ReduceLROnPlateau(optimizer, mode=\"min\", factor=0.5, patience=2, verbose=True)","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:47:11.058688Z","iopub.execute_input":"2025-12-09T18:47:11.058953Z","iopub.status.idle":"2025-12-09T18:47:11.082001Z","shell.execute_reply.started":"2025-12-09T18:47:11.058934Z","shell.execute_reply":"2025-12-09T18:47:11.081398Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 7. Training Loop\nprint(\"\\n🚀 Training...\")\nbest_val_loss = float('inf')\npatience_counter = 0\nbest_model_path = WORK_DIR / \"best_model.pt\"\n\nhistory = {\n    \"train_loss\": [],\n    \"val_loss\": [],\n    \"val_f1_bp\": [],\n    \"val_f1_mf\": [],\n    \"val_f1_cc\": [],\n}\n\nfor epoch in range(EPOCHS):\n    train_loss = train_epoch(\n        model,\n        train_loader,\n        criterion_bp,\n        criterion_mf,\n        criterion_cc,\n        optimizer,\n        DEVICE,\n        w_bp=W_BP,\n        w_mf=W_MF,\n        w_cc=W_CC,\n    )\n\n    (\n        val_loss,\n        val_bp_preds,\n        val_mf_preds,\n        val_cc_preds,\n        val_bp_labels,\n        val_mf_labels,\n        val_cc_labels,\n    ) = evaluate(\n        model,\n        val_loader,\n        criterion_bp,\n        criterion_mf,\n        criterion_cc,\n        DEVICE,\n        w_bp=W_BP,\n        w_mf=W_MF,\n        w_cc=W_CC,\n    )\n\n    # metric từng ontology\n    val_bp_bin = (val_bp_preds >= 0.5).astype(int)\n    val_mf_bin = (val_mf_preds >= 0.5).astype(int)\n    val_cc_bin = (val_cc_preds >= 0.5).astype(int)\n\n    # 1) Metric không IA: debug xem model học được chưa\n    bp_unw = compute_metrics(val_bp_labels, val_bp_bin, weights=None)\n    mf_unw = compute_metrics(val_mf_labels, val_mf_bin, weights=None)\n    cc_unw = compute_metrics(val_cc_labels, val_cc_bin, weights=None)\n\n    # 2) Metric có IA: gần với metric thật của cuộc thi\n    bp_ia = compute_metrics(val_bp_labels, val_bp_bin, ia_weights_bp)\n    mf_ia = compute_metrics(val_mf_labels, val_mf_bin, ia_weights_mf)\n    cc_ia = compute_metrics(val_cc_labels, val_cc_bin, ia_weights_cc)\n\n    scheduler.step(val_loss)\n\n    history[\"train_loss\"].append(train_loss)\n    history[\"val_loss\"].append(val_loss)\n    history[\"val_f1_bp\"].append(bp_unw[\"f1\"])\n    history[\"val_f1_mf\"].append(mf_unw[\"f1\"])\n    history[\"val_f1_cc\"].append(cc_unw[\"f1\"])\n\n    print(f\"Epoch {epoch+1}/{EPOCHS}\")\n    print(f\"  Train Loss: {train_loss:.4f}\")\n    print(f\"  Val Loss:   {val_loss:.4f}\")\n\n    # log metric không IA: dễ đọc, phục vụ debug (theo lời người hướng dẫn)\n    print(\n        \"  [Unweighted] \"\n        f\"BP P/R/F1: {bp_unw['precision']:.4f}/{bp_unw['recall']:.4f}/{bp_unw['f1']:.4f} | \"\n        f\"MF P/R/F1: {mf_unw['precision']:.4f}/{mf_unw['recall']:.4f}/{mf_unw['f1']:.4f} | \"\n        f\"CC P/R/F1: {cc_unw['precision']:.4f}/{cc_unw['recall']:.4f}/{cc_unw['f1']:.4f}\"\n    )\n\n    # log thêm IA F1 để tham chiếu, chấp nhận rất nhỏ\n    print(\n        \"  [IA-weighted F1] \"\n        f\"BP: {bp_ia['f1']:.6f}  MF: {mf_ia['f1']:.6f}  CC: {cc_ia['f1']:.6f}\"\n    )\n\n\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        patience_counter = 0\n        torch.save(model.state_dict(), best_model_path)\n        print(\"  ✓ Saved best model\")\n    else:\n        patience_counter += 1\n        if patience_counter >= PATIENCE:\n            print(f\"\\n  Early stopping triggered after {epoch+1} epochs\")\n            break\n\n    print()\n\nprint(\"Loading best model...\")\nmodel.load_state_dict(torch.load(best_model_path))","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:46:52.671879Z","iopub.execute_input":"2025-12-09T18:46:52.672165Z","iopub.status.idle":"2025-12-09T18:46:52.689856Z","shell.execute_reply.started":"2025-12-09T18:46:52.672136Z","shell.execute_reply":"2025-12-09T18:46:52.688970Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 8. Find Best Threshold\n# ----------------------------------------------------------------------------\nprint(\"\\n📈 Evaluating on validation set (multi-task)...\")\n(\n    _,\n    val_bp_preds,\n    val_mf_preds,\n    val_cc_preds,\n    val_bp_labels,\n    val_mf_labels,\n    val_cc_labels,\n) = evaluate(\n    model,\n    val_loader,\n    criterion_bp,\n    criterion_mf,\n    criterion_cc,\n    DEVICE,\n    w_bp=W_BP,\n    w_mf=W_MF,\n    w_cc=W_CC,\n)\n\nif THRESHOLD_SEARCH:\n    # IA-weighted\n    best_thr_bp_ia, best_f1_bp_ia = find_best_threshold(\n        val_bp_labels, val_bp_preds, ia_weights_bp, THRESHOLD_GRID\n    )\n    # non-IA\n    best_thr_bp_unw, best_f1_bp_unw = find_best_threshold(\n        val_bp_labels, val_bp_preds, None, THRESHOLD_GRID\n    )\n    print(\n        f\"BP threshold IA: {best_thr_bp_ia:.2f}, F1_IA: {best_f1_bp_ia:.6f} | \"\n        f\"BP threshold unweighted: {best_thr_bp_unw:.2f}, F1_unw: {best_f1_bp_unw:.4f}\"\n    )","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:47:14.018606Z","iopub.execute_input":"2025-12-09T18:47:14.019188Z","iopub.status.idle":"2025-12-09T18:47:14.033676Z","shell.execute_reply.started":"2025-12-09T18:47:14.019165Z","shell.execute_reply":"2025-12-09T18:47:14.032880Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 9. Generate Test Predictions\nprint(\"\\n🔮 Generating test predictions...\")\n\n# Load test embeddings\ntest_emb = np.load(PROTBERT_TEST_EMB).astype(np.float32)\nprint(f\"  Test embeddings shape: {test_emb.shape}\")\n\n# Create test dataset\ntest_dataset = ProteinDataset(\n    test_emb,\n    np.zeros((len(test_emb), y_bp.shape[1])),\n    np.zeros((len(test_emb), y_mf.shape[1])),\n    np.zeros((len(test_emb), y_cc.shape[1])),\n)\ntest_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE * 2, \n                        shuffle=False, num_workers=2, pin_memory=True)\n\n# Predict\nmodel.eval()\nall_bp_preds, all_mf_preds, all_cc_preds = [], [], []\n\nwith torch.no_grad():\n    for embeddings, _, _, _ in tqdm(test_loader, desc=\"Predicting\"):\n        embeddings = embeddings.to(DEVICE)\n        logits_bp, logits_mf, logits_cc = model(embeddings)\n\n        prob_bp = torch.sigmoid(logits_bp)\n        prob_mf = torch.sigmoid(logits_mf)\n        prob_cc = torch.sigmoid(logits_cc)\n\n        all_bp_preds.append(prob_bp.cpu().numpy())\n        all_mf_preds.append(prob_mf.cpu().numpy())\n        all_cc_preds.append(prob_cc.cpu().numpy())\n\ntest_bp_preds = np.vstack(all_bp_preds)\ntest_mf_preds = np.vstack(all_mf_preds)\ntest_cc_preds = np.vstack(all_cc_preds)\n\nprint(f\"  BP preds shape: {test_bp_preds.shape}\")\nprint(f\"  MF preds shape: {test_mf_preds.shape}\")\nprint(f\"  CC preds shape: {test_cc_preds.shape}\")","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:42:39.543726Z","iopub.execute_input":"2025-12-09T18:42:39.543927Z","iopub.status.idle":"2025-12-09T18:43:03.007395Z","shell.execute_reply.started":"2025-12-09T18:42:39.543909Z","shell.execute_reply":"2025-12-09T18:43:03.006388Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 10. Propagate Predictions\nif PROPAGATE_PRED:\n    test_bp_preds = propagate_predictions(\n        test_bp_preds, parents_map, list(mlb_bp.classes_), PROPAGATE_ITERATIONS\n    )\n    test_mf_preds = propagate_predictions(\n        test_mf_preds, parents_map, list(mlb_mf.classes_), PROPAGATE_ITERATIONS\n    )\n    test_cc_preds = propagate_predictions(\n        test_cc_preds, parents_map, list(mlb_cc.classes_), PROPAGATE_ITERATIONS\n    )\n\n# concat theo thứ tự BP -> MF -> CC\nall_test_preds = np.concatenate([test_bp_preds, test_mf_preds, test_cc_preds], axis=1)\nall_terms = np.concatenate([mlb_bp.classes_, mlb_mf.classes_, mlb_cc.classes_])\ntest_ids = list(test_seqs.keys())[:all_test_preds.shape[0]]","metadata":{"execution":{"iopub.status.busy":"2025-12-09T18:43:03.008963Z","iopub.execute_input":"2025-12-09T18:43:03.009247Z","iopub.status.idle":"2025-12-09T18:43:42.525479Z","shell.execute_reply.started":"2025-12-09T18:43:03.009216Z","shell.execute_reply":"2025-12-09T18:43:42.524697Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 10b. Save scores for ensemble\nSCORES_PROTMT = WORK_DIR / \"scores_protmt.npy\"\nIDS_PROTMT    = WORK_DIR / \"test_ids_protmt.npy\"\nCLASSES_PROTMT= WORK_DIR / \"classes_protmt.npy\"\n\nnp.save(SCORES_PROTMT, all_test_preds.astype(np.float16))\nnp.save(IDS_PROTMT,    np.array(test_ids))\nnp.save(CLASSES_PROTMT,np.array(all_terms))\n\nprint(f\"✅ Saved multi-task ProtBERT scores to {SCORES_PROTMT}\")\nprint(f\"✅ Saved test IDs to {IDS_PROTMT}\")\nprint(f\"✅ Saved classes to {CLASSES_PROTMT}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 11. Write Submission\nprint(f\"\\n💾 Writing submission to {OUTPUT_FILE}...\")\n\nwith open(OUTPUT_FILE, \"w\") as f:\n    for i, protein_id in enumerate(tqdm(test_ids, desc=\"Writing\")):\n        probs = all_test_preds[i]\n        top_indices = np.argsort(probs)[-TOP_K_PER_PROTEIN:][::-1]\n    \n        for idx in top_indices:\n            score = float(probs[idx])\n            if score > 1e-6:\n                go_term = all_terms[idx]\n                f.write(f\"{protein_id}\\t{go_term}\\t{score:.3f}\\n\")\n\n\nprint(\"\\n\" + \"=\"*70)\nprint(\"✅ PIPELINE COMPLETED SUCCESSFULLY\")\nprint(\"=\"*70)\n\n# Clean up\ndel model, train_loader, val_loader, test_loader\ntorch.cuda.empty_cache()\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T18:44:20.250733Z","iopub.execute_input":"2025-12-09T18:44:20.251441Z","iopub.status.idle":"2025-12-09T18:45:36.386745Z","shell.execute_reply.started":"2025-12-09T18:44:20.251413Z","shell.execute_reply":"2025-12-09T18:45:36.385664Z"}},"outputs":[],"execution_count":null}]}