{"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":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14066577,"sourceType":"datasetVersion","datasetId":8953533},{"sourceId":14066843,"sourceType":"datasetVersion","datasetId":8950709}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install protobuf==4.25.3","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:20.302045Z","iopub.execute_input":"2025-12-12T18:09:20.302387Z","iopub.status.idle":"2025-12-12T18:09:23.386540Z","shell.execute_reply.started":"2025-12-12T18:09:20.302362Z","shell.execute_reply":"2025-12-12T18:09:23.385561Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, gc\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nfrom typing import Dict, List, Set, Tuple, Optional\n\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.feature_extraction.text import TfidfVectorizer\nfrom sklearn.linear_model import SGDClassifier, LogisticRegression\nfrom sklearn.multiclass import OneVsRestClassifier\nimport random\nimport joblib\nimport scipy.sparse as sp","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.387979Z","iopub.execute_input":"2025-12-12T18:09:23.388366Z","iopub.status.idle":"2025-12-12T18:09:23.393769Z","shell.execute_reply.started":"2025-12-12T18:09:23.388341Z","shell.execute_reply":"2025-12-12T18:09:23.392999Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Config:\n    \"\"\"Centralized configuration\"\"\"\n    # Paths\n    DATA_DIR = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\n    TRAIN_DIR = DATA_DIR / \"Train\"\n    TEST_DIR = DATA_DIR / \"Test\"\n    WORK_DIR = Path(\"/kaggle/working\")\n    \n    # Data files\n    TRAIN_FASTA = TRAIN_DIR / \"train_sequences.fasta\"\n    TRAIN_TERMS = TRAIN_DIR / \"train_terms.tsv\"\n    TRAIN_TAXONOMY = TRAIN_DIR / \"train_taxonomy.tsv\"\n    GO_OBO = TRAIN_DIR / \"go-basic.obo\"\n    IA_FILE = DATA_DIR / \"IA.tsv\"\n    TEST_FASTA = TEST_DIR / \"testsuperset.fasta\"\n    SAMPLE_SUBMISSION = DATA_DIR / \"sample_submission.tsv\"\n    OUTPUT_FILE = WORK_DIR / \"submission.tsv\"\n    \n    # NEW: Embedding = K-mer TF-IDF cho protein\n    # K = 3 là lựa chọn hợp lý cho protein (20 aa → 20^3 ≈ 8000 3-mer).\n    # Giới hạn số feature để tránh OOM.\n    \n    # TF-IDF k-mer\n    KMER_NGRAM_RANGE = (3, 3)\n    KMER_MAX_FEATURES = 50000\n    KMER_MIN_DF = 2\n    KMER_SUBLINEAR_TF = True      # sublinear_tf thường giúp\n\n    MODEL_TYPE = \"sgd\"            # \"sgd\" hoặc \"logreg\"\n    \n    # SGD (khuyến nghị cho nhanh + scale tốt)\n    SGD_ALPHA = 2e-6              # 1e-6 ~ 5e-6\n    SGD_MAX_ITER = 30             # 20~50\n    SGD_TOL = 1e-3\n    SGD_N_JOBS = -1\n    \n    # Logistic Regression (chậm hơn, thử nếu có thời gian)\n    LOGREG_C = 4.0\n    LOGREG_MAX_ITER = 200\n    LOGREG_N_JOBS = -1\n\n    # Model parameters (đã chỉnh lại cho CAFA + TF-IDF)\n    RANDOM_SEED = 42\n    TOP_K_LABELS = 3000        # như cũ, top 3000 GO term\n    HIDDEN_UNITS = [1024, 512] # MLP to hơn vì input TF-IDF thấp hơn ProtBERT\n    DROPOUT = 0.5              # regularization mạnh hơn\n    LEARNING_RATE = 3e-4       # LR nhỏ hơn một chút cho ổn định\n    BATCH_SIZE = 64            # nếu OOM có thể hạ xuống 32\n    EPOCHS = 50               # EarlyStopping sẽ cắt sớm\n    PATIENCE = 5\n    \n    # Prediction parameters\n    TOP_K_PER_PROTEIN = 200\n    \n    # BEST THRESHOLD với CAFA thường rất thấp ~0.01–0.1\n    THRESHOLD_SEARCH = True\n    THRESHOLD_GRID = np.concatenate([\n        np.arange(0.001, 0.051, 0.002),\n        np.arange(0.05, 0.201, 0.01)\n    ])\n        \n    # GO propagation\n    PROPAGATE_TRAIN = True\n    PROPAGATE_PRED = True\n    PROPAGATE_ITERATIONS = 3\n    SGD_N_JOBS = 4\n    \n    @classmethod\n    def set_seed(cls):\n        np.random.seed(cls.RANDOM_SEED)\n        random.seed(cls.RANDOM_SEED)\n        os.environ[\"PYTHONHASHSEED\"] = str(cls.RANDOM_SEED)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.394717Z","iopub.execute_input":"2025-12-12T18:09:23.395035Z","iopub.status.idle":"2025-12-12T18:09:23.408216Z","shell.execute_reply.started":"2025-12-12T18:09:23.395012Z","shell.execute_reply":"2025-12-12T18:09:23.407434Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import time, os\n\ntry:\n    import psutil\n    _PROC = psutil.Process(os.getpid())\nexcept Exception:\n    _PROC = None\n\nclass Timer:\n    def __init__(self):\n        self.t = time.perf_counter()\n\n    def hit(self, msg: str):\n        now = time.perf_counter()\n        dt = now - self.t\n        self.t = now\n        if _PROC is not None:\n            rss_gb = _PROC.memory_info().rss / (1024**3)\n            print(f\"[TIMER] {msg}: {dt:.2f}s | RSS={rss_gb:.2f} GB\")\n        else:\n            print(f\"[TIMER] {msg}: {dt:.2f}s\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.409911Z","iopub.execute_input":"2025-12-12T18:09:23.410128Z","iopub.status.idle":"2025-12-12T18:09:23.422268Z","shell.execute_reply.started":"2025-12-12T18:09:23.410112Z","shell.execute_reply":"2025-12-12T18:09:23.421649Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class DataLoader:\n    \"\"\"Handle all data loading operations\"\"\"\n    \n    @staticmethod\n    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                    # Save previous sequence\n                    if protein_id:\n                        sequences[protein_id] = \"\".join(seq_parts)\n                    \n                    # Parse new header\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            # Save last sequence\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    @staticmethod\n    def 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    @staticmethod\n    def 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":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.423158Z","iopub.execute_input":"2025-12-12T18:09:23.423389Z","iopub.status.idle":"2025-12-12T18:09:23.434889Z","shell.execute_reply.started":"2025-12-12T18:09:23.423375Z","shell.execute_reply":"2025-12-12T18:09:23.434226Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class GOGraph:\n    \"\"\"Handle Gene Ontology graph operations\"\"\"\n    \n    def __init__(self, obo_path: Path):\n        self.parents, self.children = self._parse_obo(obo_path)\n    \n    def _parse_obo(self, path: Path) -> Tuple[Dict, Dict]:\n        \"\"\"Parse OBO file to extract parent-child relationships\"\"\"\n        parents = defaultdict(set)\n        children = defaultdict(set)\n        \n        if not path.exists():\n            print(\"⚠ OBO file not found\")\n            return parents, children\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(\"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        return dict(parents), dict(children)\n    \n    def get_ancestors(self, go_term: str) -> 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 self.parents.get(current, []):\n                if parent not in ancestors:\n                    ancestors.add(parent)\n                    stack.append(parent)\n        \n        return ancestors\n    \n    def propagate_labels(self, annotations: Dict[str, List[str]]) -> 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 annotations.items():\n            expanded = set(terms)\n            for term in terms:\n                expanded.update(self.get_ancestors(term))\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":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.435485Z","iopub.execute_input":"2025-12-12T18:09:23.435697Z","iopub.status.idle":"2025-12-12T18:09:23.450560Z","shell.execute_reply.started":"2025-12-12T18:09:23.435675Z","shell.execute_reply":"2025-12-12T18:09:23.449973Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === SPEED FIX: cache ancestors ===\nfrom functools import lru_cache\n\ndef _attach_cached_ancestors(go_graph: GOGraph):\n    @lru_cache(maxsize=None)\n    def ancestors_cached(term: str):\n        ancestors = set()\n        stack = [term]\n        while stack:\n            cur = stack.pop()\n            for p in go_graph.parents.get(cur, ()):\n                if p not in ancestors:\n                    ancestors.add(p)\n                    stack.append(p)\n        # trả về tuple để cache ổn định\n        return tuple(ancestors)\n\n    go_graph.get_ancestors = lambda t: set(ancestors_cached(t))  # override method\n\n# dùng sau khi tạo go_graph\n# pipeline.go_graph = GOGraph(...)\n# _attach_cached_ancestors(pipeline.go_graph)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.451552Z","iopub.execute_input":"2025-12-12T18:09:23.451812Z","iopub.status.idle":"2025-12-12T18:09:23.464465Z","shell.execute_reply.started":"2025-12-12T18:09:23.451797Z","shell.execute_reply":"2025-12-12T18:09:23.463860Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class EmbeddingHandler:\n    \"\"\"Handle protein embeddings: K-mer TF-IDF (SPARSE)\"\"\"\n\n    @staticmethod\n    def build_kmer_tfidf(\n        sequences: List[str],\n        ngram_range: Tuple[int, int],\n        max_features: int,\n        min_df: int = 1,\n        sublinear_tf: bool = True,\n    ):\n        print(f\"→ Building TF-IDF ngram_range={ngram_range} (max_features={max_features}, min_df={min_df})...\")\n\n        vectorizer = TfidfVectorizer(\n            analyzer=\"char\",\n            ngram_range=ngram_range,\n            lowercase=False,\n            max_features=max_features,\n            min_df=min_df,\n            sublinear_tf=sublinear_tf,\n            dtype=np.float32,\n            norm=\"l2\",\n            use_idf=True,\n        )\n        X = vectorizer.fit_transform(sequences)  # SPARSE\n        print(f\"✓ Built TF-IDF: X={X.shape}, vocab_size={len(vectorizer.vocabulary_):,}\")\n        return X, vectorizer\n\n    @staticmethod\n    def transform_kmer_tfidf(\n        sequences: List[str],\n        vectorizer: TfidfVectorizer,\n    ):\n        X = vectorizer.transform(sequences)      # SPARSE\n        print(f\"✓ Transformed sequences with TF-IDF: {X.shape}\")\n        return X\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.465087Z","iopub.execute_input":"2025-12-12T18:09:23.465384Z","iopub.status.idle":"2025-12-12T18:09:23.476797Z","shell.execute_reply.started":"2025-12-12T18:09:23.465363Z","shell.execute_reply":"2025-12-12T18:09:23.476155Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ModelBuilder:\n    \"\"\"Build sklearn linear multi-label models\"\"\"\n\n    @staticmethod\n    def build_linear_model(config: Config):\n        if config.MODEL_TYPE == \"logreg\":\n            base = LogisticRegression(\n                C=config.LOGREG_C,\n                solver=\"saga\",\n                max_iter=config.LOGREG_MAX_ITER,\n                n_jobs=config.LOGREG_N_JOBS,\n                verbose=0,\n            )\n            model = OneVsRestClassifier(base, n_jobs=config.LOGREG_N_JOBS)\n            return model\n\n        # default: SGD\n        base = SGDClassifier(\n            loss=\"log_loss\",\n            alpha=config.SGD_ALPHA,\n            max_iter=config.SGD_MAX_ITER,\n            tol=config.SGD_TOL,\n            early_stopping=False,\n            n_iter_no_change=3,\n            validation_fraction=0.1,\n            average=True,\n        )\n        model = OneVsRestClassifier(base, n_jobs=config.SGD_N_JOBS)\n        return model\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.477516Z","iopub.execute_input":"2025-12-12T18:09:23.477738Z","iopub.status.idle":"2025-12-12T18:09:23.490987Z","shell.execute_reply.started":"2025-12-12T18:09:23.477724Z","shell.execute_reply":"2025-12-12T18:09:23.490287Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class 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        # y_true có thể là CSR sparse (do sparse_output=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)      # <-- FIX\n        best_t, best_f = 0.5, -1.0\n\n        for t in thresholds:\n            y_pred_b = (y_prob >= t)\n\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)      # <-- FIX\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-12T18:09:23.493003Z","iopub.execute_input":"2025-12-12T18:09:23.493628Z","iopub.status.idle":"2025-12-12T18:09:23.505001Z","shell.execute_reply.started":"2025-12-12T18:09:23.493610Z","shell.execute_reply":"2025-12-12T18:09:23.504227Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class PredictionPropagator:\n    \"\"\"Propagate predictions up GO hierarchy\"\"\"\n    \n    def __init__(self, go_graph: GOGraph, go_terms: List[str]):\n        self.go_terms = go_terms\n        self.term_to_idx = {t: i for i, t in enumerate(go_terms)}\n        \n        # Restrict parents to terms in our vocabulary\n        self.parents = {}\n        for term in go_terms:\n            self.parents[term] = {p for p in go_graph.parents.get(term, []) \n                                 if p in self.term_to_idx}\n    \n    def propagate(self, predictions: np.ndarray, \n                  iterations: int = 3) -> np.ndarray:\n        \"\"\"Propagate predictions iteratively\"\"\"\n        pred_copy = predictions.copy()\n        \n        for _ in range(iterations):\n            changed = False\n            \n            for child_idx, child_term in enumerate(self.go_terms):\n                child_scores = pred_copy[:, child_idx]\n                \n                for parent_term in self.parents.get(child_term, []):\n                    parent_idx = self.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                break\n        \n        return pred_copy","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.505738Z","iopub.execute_input":"2025-12-12T18:09:23.505930Z","iopub.status.idle":"2025-12-12T18:09:23.519021Z","shell.execute_reply.started":"2025-12-12T18:09:23.505910Z","shell.execute_reply":"2025-12-12T18:09:23.518435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CAFA6Pipeline:\n    \"\"\"Main training and prediction pipeline\"\"\"\n    \n    def __init__(self, config: Config):\n        self.config = config\n        config.set_seed()\n        \n        # Initialize components\n        print(\"\\n\" + \"=\"*70)\n        print(\"CAFA-6 PROTEIN FUNCTION PREDICTION PIPELINE\")\n        print(\"=\"*70 + \"\\n\")\n        \n        print(\"📁 Loading data...\")\n        self.loader = DataLoader()\n        self.train_seqs = self.loader.read_fasta(config.TRAIN_FASTA)\n        self.test_seqs = self.loader.read_fasta(config.TEST_FASTA)\n        self.annotations = self.loader.read_annotations(config.TRAIN_TERMS)\n        self.ia_weights = self.loader.read_ia_weights(config.IA_FILE)\n        \n        print(\"\\n🌳 Loading GO ontology...\")\n        self.go_graph = GOGraph(config.GO_OBO)\n        _attach_cached_ancestors(self.go_graph)\n        \n        # Propagate training labels\n        if config.PROPAGATE_TRAIN:\n            self.annotations = self.go_graph.propagate_labels(self.annotations)\n        \n        # Prepare labels\n        self._prepare_labels()\n        \n        print(\"\\n🧬 Preparing K-mer TF-IDF embeddings...\")\n        self.embedding_handler = EmbeddingHandler()\n        self.vectorizer = None  \n        \n    def _prepare_labels(self):\n        \"\"\"Prepare label matrix\"\"\"\n        print(\"\\n🏷️  Preparing labels...\")\n        \n        # Get proteins with both sequences and annotations\n        self.train_proteins = [p for p in self.annotations.keys() \n                              if p in self.train_seqs]\n        print(f\"  {len(self.train_proteins):,} training proteins\")\n        \n        # Count term frequencies\n        term_counts = Counter()\n        for protein in self.train_proteins:\n            term_counts.update(self.annotations[protein])\n        \n        # Select top-K terms\n        top_terms = [t for t, _ in term_counts.most_common(self.config.TOP_K_LABELS)]\n        self.chosen_terms = set(top_terms)\n        print(f\"  Using top {len(self.chosen_terms):,} GO terms\")\n        \n        # Filter annotations\n        for protein in self.train_proteins:\n            self.annotations[protein] = [t for t in self.annotations[protein] \n                                        if t in self.chosen_terms]\n        \n        # Create label matrix\n        labels_list = [self.annotations[p] for p in self.train_proteins]\n        self.mlb = MultiLabelBinarizer(classes=sorted(self.chosen_terms), sparse_output=True)\n        self.y = self.mlb.fit_transform(labels_list)  # CSR sparse\n        \n        print(f\"  Label matrix shape: {self.y.shape}\")\n    \n    def prepare_data(self):\n        \"\"\"Build K-mer TF-IDF embeddings cho training\"\"\"\n        print(\"\\n📊 Preparing training data (K-mer TF-IDF)...\")\n        t = Timer()\n        \n        # Lấy chuỗi train theo đúng thứ tự self.train_proteins\n        train_texts = [self.train_seqs[p] for p in self.train_proteins]\n        \n        # Build TF-IDF\n        self.X, self.vectorizer = self.embedding_handler.build_kmer_tfidf(\n            train_texts,\n            ngram_range=self.config.KMER_NGRAM_RANGE,\n            max_features=self.config.KMER_MAX_FEATURES,\n            min_df=self.config.KMER_MIN_DF,\n            sublinear_tf=self.config.KMER_SUBLINEAR_TF,\n        )\n\n        \n        # Train/val split\n        self.X_train, self.X_val, self.y_train, self.y_val = train_test_split(\n            self.X, self.y, test_size=0.15, random_state=self.config.RANDOM_SEED\n        )\n        t.hit(\"TF-IDF fit_transform\")\n        \n        print(f\"  Train: {self.X_train.shape}\")\n        print(f\"  Val:   {self.X_val.shape}\")\n        \n    def build_and_train(self):\n        print(\"\\n🏗️  Building linear model (OneVsRest)...\")\n        self.model = ModelBuilder.build_linear_model(self.config)\n        \n        t = Timer()\n        print(\"🚀 Training (sklearn)...\")\n        with joblib.parallel_backend(\"threading\"):\n            self.model.fit(self.X_train, self.y_train)\n\n        t.hit(\"model.fit\")\n        \n    def evaluate(self):\n        print(\"\\n📈 Evaluating (Fmax sweep)...\")\n    \n        self.evaluator = Evaluator(self.ia_weights, list(self.mlb.classes_))\n    \n        # predict_proba for OneVsRest -> shape (n_samples, n_classes)\n        y_val_prob = self.model.predict_proba(self.X_val).astype(np.float32)\n    \n        # Có thể propagate trên VAL để threshold phản ánh đúng pipeline submit\n        if self.config.PROPAGATE_PRED:\n            print(\"  Propagating VAL predictions...\")\n            propagator = PredictionPropagator(self.go_graph, list(self.mlb.classes_))\n            y_val_prob = propagator.propagate(y_val_prob, self.config.PROPAGATE_ITERATIONS)\n    \n        if self.config.THRESHOLD_SEARCH:\n            t_best, f_best = self.evaluator.fmax(self.y_val, y_val_prob, self.config.THRESHOLD_GRID)\n            t_best_ia, f_best_ia = self.evaluator.ia_fmax(self.y_val, y_val_prob, self.config.THRESHOLD_GRID)\n    \n            self.best_threshold = t_best\n            self.best_f1 = f_best\n            print(f\"  Best threshold (Fmax): {t_best:.4f} | Fmax: {f_best:.4f}\")\n            print(f\"  Best threshold (IA-Fmax): {t_best_ia:.4f} | IA-Fmax: {f_best_ia:.4f}\")\n        else:\n            self.best_threshold = 0.05\n            print(f\"  Using fixed threshold: {self.best_threshold:.4f}\")\n\n    def predict_and_submit(self):\n        print(\"\\n🔮 Generating predictions (batch + streaming write)...\")\n    \n        if self.vectorizer is None:\n            raise RuntimeError(\"TF-IDF vectorizer is None. Did you call prepare_data() before predict_and_submit()?\")\n    \n        test_ids = list(self.test_seqs.keys())\n        test_texts = [self.test_seqs[p] for p in test_ids]\n    \n        propagator = None\n        if self.config.PROPAGATE_PRED:\n            propagator = PredictionPropagator(self.go_graph, list(self.mlb.classes_))\n    \n        top_k = self.config.TOP_K_PER_PROTEIN\n        batch_size = 4096  # chỉnh 2048 nếu RAM căng\n    \n        print(f\"\\n💾 Writing submission to {self.config.OUTPUT_FILE}...\")\n        with open(self.config.OUTPUT_FILE, \"w\") as f:\n            for start in range(0, len(test_ids), batch_size):\n                end = min(start + batch_size, len(test_ids))\n                batch_ids = test_ids[start:end]\n                batch_texts = test_texts[start:end]\n    \n                X_batch = self.embedding_handler.transform_kmer_tfidf(batch_texts, self.vectorizer)\n                y_prob = self.model.predict_proba(X_batch).astype(np.float32)\n    \n                if propagator is not None:\n                    y_prob = propagator.propagate(y_prob, self.config.PROPAGATE_ITERATIONS)\n    \n                # write topK\n                for i, protein_id in enumerate(batch_ids):\n                    probs = y_prob[i]\n    \n                    # topK nhanh hơn argsort full\n                    if top_k < probs.shape[0]:\n                        idx = np.argpartition(probs, -top_k)[-top_k:]\n                        idx = idx[np.argsort(probs[idx])[::-1]]\n                    else:\n                        idx = np.argsort(probs)[::-1]\n    \n                    for j in idx:\n                        score = float(probs[j])\n                        if score > 1e-8:\n                            go_term = self.mlb.classes_[j]\n                            f.write(f\"{protein_id}\\t{go_term}\\t{score:.6f}\\n\")\n    \n        print(\"✅ Done!\")\n\n        \n    def run(self):\n        t = Timer()\n        try:\n            t.hit(\"start run()\")\n    \n            self.prepare_data()\n            t.hit(\"prepare_data() done\")\n    \n            self.build_and_train()\n            t.hit(\"build_and_train() done\")\n    \n            self.evaluate()\n            t.hit(\"evaluate() done\")\n    \n            self.predict_and_submit()\n            t.hit(\"predict_and_submit() done\")\n    \n            print(\"\\n\" + \"=\"*70)\n            print(\"PIPELINE COMPLETED SUCCESSFULLY\")\n            print(\"=\"*70)\n    \n        except Exception as e:\n            print(f\"\\n❌ Error: {e}\")\n            raise\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.519705Z","iopub.execute_input":"2025-12-12T18:09:23.519924Z","iopub.status.idle":"2025-12-12T18:09:23.540404Z","shell.execute_reply.started":"2025-12-12T18:09:23.519908Z","shell.execute_reply":"2025-12-12T18:09:23.539778Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize and run pipeline\nconfig = Config()\npipeline = CAFA6Pipeline(config)\npipeline.run()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T18:09:23.541170Z","iopub.execute_input":"2025-12-12T18:09:23.541408Z","iopub.status.idle":"2025-12-12T18:11:12.971266Z","shell.execute_reply.started":"2025-12-12T18:09:23.541388Z","shell.execute_reply":"2025-12-12T18:11:12.969998Z"}},"outputs":[],"execution_count":null}]}