{"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":13680623,"sourceType":"datasetVersion","datasetId":8699749},{"sourceId":14153916,"sourceType":"datasetVersion","datasetId":9021225},{"sourceId":14154204,"sourceType":"datasetVersion","datasetId":9021432},{"sourceId":14154618,"sourceType":"datasetVersion","datasetId":9021718}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install BIO","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#Sample submission\nimport pandas as pd\nsub = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/sample_submission.tsv\", sep= \"\\t\", on_bad_lines='skip', header = None)\n#sub.columns = [\"The Protein ID\", \"The Gene Ontology term (GO) ID\", \"Predicted link probability that GO appear in Protein\"]\nsub.head(5)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"MAIN_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\nclass config:\n    train_sequences_path = MAIN_DIR  + \"/Train/train_sequences.fasta\"\n    train_term_path = MAIN_DIR +\"/Train/train_terms.tsv\"\n    train_labels_path = MAIN_DIR + \"/Train/train_terms.tsv\"\n    test_sequences_path = MAIN_DIR + \"/Test/testsuperset.fasta\"\n    taxonomy_path = os.path.join(MAIN_DIR, \"Train/train_taxonomy.tsv\")\n    n_epochs = 50\n    batch_size = 128\n    lr = 1e-5\n    early_stopping = 7\n    obo_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\n    goa_path = \"/kaggle/input/protein-go-annotations/goa_uniprot_all.csv\"\n    device = torch.device('cpu')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Utilities","metadata":{}},{"cell_type":"code","source":"from Bio import SeqIO\ndef get_train_test_df():\n    train_df = pd.read_csv(config.train_term_path, sep=\"\\t\")\n    test_df = pd.read_csv(config.test_sequences_path, sep = \"\\t\")\n    train_ids = train_df[\"EntryID\"].unique().tolist()\n    test_ids = []\n    for record in SeqIO.parse(config.test_sequences_path, \"fasta\"):\n        test_id = record.id.split()[0]\n        test_ids.append(test_id)\n    return train_df, train_ids, test_df, test_ids\ndef preprocess_data(model_name = \"esm2\"):\n    train_df, train_ids, test_df, test_ids = get_train_test_df()\n    if model_name == \"esm2\":\n        ids = pd.read_csv(\"/kaggle/input/cafa6-protein-embeddings-esm2/protein_ids.csv\")[\"protein_id\"].tolist()\n        id_to_index = {pid: i for i, pid in enumerate(ids)}\n        #Converting to data to train , because the embedding is already in order for each protein\n        train_idx = [id_to_index[i] for i in train_ids if i in id_to_index]\n        test_idx  = [id_to_index[i] for i in test_ids  if i in id_to_index]\n        #init train and test\n        emb = np.load(\"/kaggle/input/cafa6-protein-embeddings-esm2/protein_embeddings.npy\")\n        X_train = emb[train_idx]\n        X_test  = emb[test_idx] \n    elif model_name ==\"t5\":\n        # train_ids_unprocessed = np.load(\"/kaggle/input/cafa-6-t5-embeddings/train_ids.npy\")\n        # id_to_index = {pid: i for i, pid in enumerate(ids)}\n        \n        # train_ids = [protein_id.split('|')[1] for protein_id in train_ids_unprocessed]\n        \n        # train_idx = [id_to_index[i] for i in train_ids if i in id_to_index]\n        # test_idx  = [id_to_index[i] for i in test_ids  if i in id_to_index]\n        train_ids = np.load(\"/kaggle/input/cafa-6-t5-embeddings/train_ids.npy\")\n        train_ids = np.load(\"/kaggle/input/cafa-6-t5-embeddings/train_ids.npy\")\n        test_ids = np.load(\"/kaggle/input/cafa-6-t5-embeddings/test_ids.npy\")\n        train_ids = [train_id.split(\"|\")[1] for train_id in train_ids]\n        id_to_index = {pid: i for i, pid in enumerate(train_ids)}\n        train_df[\"order\"] = train_df[\"EntryID\"].map(id_to_index)\n        train_df = train_df.sort_values(\"order\").reset_index(drop=True)\n        test_ids = list(np.load(\"/kaggle/input/cafa-6-t5-embeddings/test_ids.npy\"))\n\n        X_train = np.load(\"/kaggle/input/cafa-6-t5-embeddings/train_embeds.npy\")\n        X_test  = np.load(\"/kaggle/input/cafa-6-t5-embeddings/test_embeds.npy\")\n        \n    np.save(\"X_train.npy\", X_train)\n    np.save(\"X_test.npy\", X_test)\n    pd.DataFrame(train_ids).to_csv(\"y_train_ids.csv\", index=False)\n    pd.DataFrame(test_ids).to_csv(\"y_test_ids.csv\", index=False)\n    # pd.DataFrame(train_ids).to_csv(\"y_train_ids.csv\", index=False)\n    # pd.DataFrame(test_ids).to_csv(\"y_test_ids.csv\", index=False)\n    from sklearn.preprocessing import MultiLabelBinarizer\n    # Group into ID and Go function for training\n    labels = train_df.groupby(\"EntryID\")[\"term\"].apply(list)\n    #Common_ids in both train_ids and labels because trains_id have unique value\n    # One hot encoder for y_train\n    Y_train = MultiLabelBinarizer().fit_transform(labels[train_ids])\n    \n    return X_train, X_test, Y_train    \n    \n#X_train, X_test, Y_train = preprocess_data(\"t5\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"class MLP(nn.Module):\n    def __init__(self, input_dim=1280, output_dim=None, \n                 NN_tuning=[1024, 512], drate=0.3, act=\"relu\"):\n        \n        super().__init__()\n\n\n        if act == \"relu\":\n            act = nn.ReLU()\n        elif act == \"gelu\":\n            act = nn.GELU()\n        elif act in [\"lrelu\", \"leakyrelu\"]:\n            act = nn.LeakyReLU()\n        else:\n            raise ValueError(f\"Unknown activation type: {regu}\")\n\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, NN_tuning[0]),\n            nn.BatchNorm1d(NN_tuning[0]),\n            act,\n            nn.Dropout(drate),\n\n            nn.Linear(NN_tuning[0], NN_tuning[1]),\n            nn.BatchNorm1d(NN_tuning[1]),\n            act,\n            nn.Dropout(drate),\n\n            nn.Linear(NN_tuning[1], output_dim)  \n        )\n\n    def forward(self, x):\n        return self.net(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class ProteinDataset(Dataset):\n    def __init__(self, X, Y):\n        # Dung de wrap vao trong pytorch\n        self.X = torch.tensor(X, dtype=torch.float32)\n        self.Y = torch.tensor(Y, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.X)\n\n    def __getitem__(self, i):\n        return self.X[i], self.Y[i]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Evaluating, predicting, and emsebmling","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\n\n# --- Fmax evaluator ---\ndef evaluate_fmax(model, val_loader, device, thresholds=None):\n    if thresholds is None:\n        thresholds = torch.linspace(0.05, 0.95, 19).to(device)\n\n    model.eval()\n\n    # For each threshold, we track TP, FP, FN\n    TP = torch.zeros(len(thresholds))\n    FP = torch.zeros(len(thresholds))\n    FN = torch.zeros(len(thresholds))\n\n    with torch.no_grad():\n        for Xb, Yb in val_loader:\n            Xb = Xb.to(device)\n\n            with torch.amp.autocast(device_type=\"cuda\"):\n                logits = model(Xb)\n\n            probs = torch.sigmoid(logits).cpu()\n            labels = Yb.cpu()\n\n            # For all thresholds at once (vectorized)\n            for i, t in enumerate(thresholds):\n                pred = (probs >= t).float()\n\n                TP[i] += (pred * labels).sum()\n                FP[i] += (pred * (1 - labels)).sum()\n                FN[i] += ((1 - pred) * labels).sum()\n\n    # Now compute F1 per threshold\n    precision = TP / (TP + FP + 1e-9)\n    recall    = TP / (TP + FN + 1e-9)\n    f1_scores = 2 * precision * recall / (precision + recall + 1e-9)\n\n    best_idx = torch.argmax(f1_scores).item()\n    best_f1 = f1_scores[best_idx].item()\n    best_thr = thresholds[best_idx].item()\n\n    return best_f1, best_thr\n\n\n\n# --- Modified val_model() ---\ndef val_model():\n    val_losses = []\n    val_f1s = []\n    val_thresholds = []\n\n    criterion = nn.BCEWithLogitsLoss()\n\n    for idx, model in enumerate(model_list):\n        model.eval()\n\n        # -------------------\n        # VALIDATION LOSS\n        # -------------------\n        val_loss = 0\n        with torch.no_grad():\n            for Xb, Yb in val_loader:\n                Xb, Yb = Xb.to(device), Yb.to(device)\n\n                with torch.amp.autocast(device_type=\"cuda\"):\n                    logits = model(Xb)\n                    loss = criterion(logits, Yb)\n\n                val_loss += loss.item() * Xb.size(0)\n\n        val_loss /= len(val_loader.dataset)\n        val_losses.append(val_loss)\n\n        # -------------------\n        # FMAX SCORE\n        # -------------------\n        fmax, thr = evaluate_fmax(model, val_loader, device)\n        val_f1s.append(fmax)\n        val_thresholds.append(thr)\n\n        print(f\"[Model {idx+1}] val_loss={val_loss:.6f} | Fmax={fmax:.4f} | thr={thr:.2f}\")\n\n        torch.cuda.empty_cache()\n\n    return val_losses, val_f1s, val_thresholds","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config.mode = \"ensemble\"\nif config.mode == \"eval\":\n    val_dataset = ProteinDataset(X_train,Y_train)\n    batch_size = config.batch_size\n    val_loader = DataLoader(val_dataset, batch_size=batch_size)\n    model_list = []\n    device = config.device\n    \n    # --- 5 ---\n    state_dict = torch.load(\"/kaggle/input/new-model-weights/model_weights_0.001243_LeakyRELU_Loss_0_3_1024_550.pt\", map_location=\"cpu\")\n    \n    model = MLP(\n        input_dim=X_train.shape[1],\n        output_dim=Y_train.shape[1],\n        NN_tuning=[1024, 550],\n        drate=0.3,\n        act=\"lrelu\"\n    ).to(device)\n    \n    model.load_state_dict(state_dict)\n    model_list.append(model)\n    val_losses, val_f1s, val_thresholds = val_model()\nelif config.mode ==\"predict\":\n    from tqdm import tqdm\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    torch.cuda.empty_cache()\n    \n    state_dict = torch.load(\"/kaggle/input/model-weights/model_weights0.00124_LOSS_0.2_(RELU_1024_512).pt\", map_location=\"cpu\")\n\n    model = MLP(\n        input_dim=X_train.shape[1],\n        output_dim=Y_train.shape[1],\n        NN_tuning=[1024, 512],\n        drate=0.3,\n        act=\"relu\"\n    ).to(device)\n    \n    model.load_state_dict(state_dict)\n    model = model.to(device)\n    model.eval()\n    model = model.half()        # reduce VRAM if needed\n          # giảm VRAM\n    \n    # --- Load test ---\n    X_test = np.load(\"/kaggle/working/X_test.npy\")\n    \n    test_ids = pd.read_csv(\"/kaggle/working/y_test_ids.csv\")[\"0\"].tolist()\n    test_id_to_index = {pid: i for i, pid in enumerate(test_ids)}\n    \n    train_terms = pd.read_csv(\n        \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n        sep=\"\\t\"\n    )\n    all_go_terms = sorted(train_terms[\"term\"].unique())\n    \n    batch_size = 32\n    min_prob = 0.02   \n    top_k = 25         \n    \n    submission_rows = []\n    \n    for i in tqdm(range(0, len(test_ids), batch_size), desc=\"Predicting & writing\"):\n        batch_ids = test_ids[i:i+batch_size]\n        batch_idx = [test_id_to_index[pid] for pid in batch_ids]\n    \n        batch = X_test[batch_idx].astype(np.float16)\n        batch = torch.tensor(batch, dtype=torch.float16).to(device)\n    \n        with torch.no_grad():\n            logits = model(batch)\n            probs = torch.sigmoid(logits).cpu().numpy()\n    \n        for j, pid in enumerate(batch_ids):\n            p = probs[j]\n    \n            idx = np.where(p >= min_prob)[0]\n    \n            if top_k is not None and len(idx) < top_k:\n                top_extra = p.argsort()[-top_k:]\n                idx = np.unique(np.concatenate([idx, top_extra]))\n    \n            # Sắp xếp theo xác suất giảm dần\n            idx = idx[np.argsort(-p[idx])]\n    \n            for k in idx:\n                term = all_go_terms[k]\n                score = float(p[k])\n                submission_rows.append([pid, term, score])\n        if len(submission_rows) > 5000:\n            df = pd.DataFrame(submission_rows, columns=[\"Id\", \"term\", \"score\"])\n            df.to_csv(\"submission.tsv\",sep=\"\\t\", mode=\"a\", header=False, index=False)\n            submission_rows = []    \n        del batch, logits, probs\n    \n        torch.cuda.empty_cache()\nelif config.mode ==\"ensemble\":\n    import pandas as pd\n    import numpy as np\n    from tqdm import tqdm\n    \n    # Ensemble weights\n    WEIGHT_MODEL_A = 0.75\n    WEIGHT_MODEL_B = 0.25\n    DIVERSITY_BONUS = 0.05\n    CONFIDENCE_THRESHOLD = 0.005\n    NOISE_LEVEL = 0.0015\n    \n    print(\"Loading Model A...\")\n    dict_a = {}\n    a_keys = set()  \n    \n    with open('/kaggle/input/best-submission/submission1.tsv', 'r') as f:\n        for line in tqdm(f, desc=\"Loading Model A\"):\n            parts = line.strip().split('\\t')\n            if len(parts) == 3:\n                protein, go_term, score = parts[0], parts[1], float(parts[2])\n                key = f\"{protein}|{go_term}\"\n                dict_a[key] = min(max(score, 0), 1.0)\n                a_keys.add(key)\n    \n    print(\"Processing ensemble...\")\n    results = []\n    \n    with open('/kaggle/input/new-best-submission/submission.tsv', 'r') as f:\n        for line in tqdm(f, desc=\"Processing Model B\"):\n            parts = line.strip().split('\\t')\n            if len(parts) == 3:\n                protein, go_term, score_b = parts[0], parts[1], float(parts[2])\n                key = f\"{protein}|{go_term}\"\n                score_a = dict_a.get(key, 0)\n                \n                # Calculate ensemble score\n                weighted_score = (score_a * WEIGHT_MODEL_A + score_b * WEIGHT_MODEL_B)\n                \n                if score_a > 0 and score_b > 0:\n                    final_score = min(weighted_score + DIVERSITY_BONUS, 1.0)\n                else:\n                    final_score = weighted_score\n                \n                # Add noise and clip\n                final_score += np.random.uniform(-NOISE_LEVEL, NOISE_LEVEL)\n                final_score = max(0, min(1, final_score))\n                \n                if final_score >= CONFIDENCE_THRESHOLD:\n                    results.append((protein, go_term, final_score))\n                \n                # Remove from dict_a to mark as processed\n                if key in dict_a:\n                    del dict_a[key]\n    \n    print(\"Processing remaining Model A entries...\")\n    for key, score_a in tqdm(dict_a.items(), desc=\"Remaining A\"):\n        protein, go_term = key.split('|')\n        score_b = 0  # Not in Model B\n        \n        weighted_score = (score_a * WEIGHT_MODEL_A + score_b * WEIGHT_MODEL_B)\n        final_score = weighted_score\n        \n        final_score += np.random.uniform(-NOISE_LEVEL, NOISE_LEVEL)\n        final_score = max(0, min(1, final_score))\n        \n        if final_score >= CONFIDENCE_THRESHOLD:\n            results.append((protein, go_term, final_score))\n    \n    del dict_a\n    import gc\n    gc.collect()\n    \n    print(f\"Saving {len(results):,} predictions...\")\n    with open('submission.tsv', 'w') as f:\n        for protein, go_term, score in tqdm(results, desc=\"Writing file\"):\n            f.write(f\"{protein}\\t{go_term}\\t{score:.6f}\\n\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Post processing","metadata":{}},{"cell_type":"code","source":"from typing import Dict, List, Set, Tuple, Optional\nimport pandas as pd\ndef parse_obo(obo_file: str) -> Tuple[Dict[str, List[str]], Dict[str, List[str]]]:\n    \"\"\"\n    Parse GO OBO file.\n    Returns: (go_parents, go_children) dictionaries.\n    \"\"\"\n    go_parents = {}   # term -> list of parent terms\n    go_children = {}  # term -> list of child terms\n    current_term = None\n    is_obsolete = False\n    \n    print(f\"Parsing GO hierarchy from {obo_file}...\")\n    \n    with open(obo_file, 'r') as f:\n        for line in f:\n            line = line.strip()\n            \n            if line == \"[Term]\":\n                current_term = None\n                is_obsolete = False\n                \n            elif line.startswith(\"id: GO:\"):\n                current_term = line.split(\"id: \")[1]\n                go_parents[current_term] = []\n                go_children.setdefault(current_term, [])\n                \n            elif line.startswith(\"is_obsolete: true\"):\n                is_obsolete = True\n                \n            elif current_term and not is_obsolete:\n                if line.startswith(\"is_a:\"):\n                    parent = line.split(\"is_a: \")[1].split(\" !\")[0].strip()\n                    go_parents[current_term].append(parent)\n                    go_children.setdefault(parent, []).append(current_term)\n                    \n                elif line.startswith(\"relationship: part_of\"):\n                    parent = line.split(\"part_of \")[1].split(\" !\")[0].strip()\n                    go_parents[current_term].append(parent)\n                    go_children.setdefault(parent, []).append(current_term)\n    \n    print(f\"Loaded {len(go_parents)} GO terms\")\n    return go_parents, go_children\n\n\ndef get_all_ancestors(term: str, go_parents: Dict, cache: Dict = None) -> Set[str]:\n    \"\"\"Get all ancestors of a GO term (with caching).\"\"\"\n    if cache is None:\n        cache = {}\n    if term in cache:\n        return cache[term]\n    \n    ancestors = set()\n    stack = [term]\n    while stack:\n        cur = stack.pop()\n        for parent in go_parents.get(cur, []):\n            if parent not in ancestors:\n                ancestors.add(parent)\n                stack.append(parent)\n    \n    cache[term] = ancestors\n    return ancestors\n\n\ndef get_all_descendants(term: str, go_children: Dict, cache: Dict = None) -> Set[str]:\n    \"\"\"Get all descendants of a GO term (with caching).\"\"\"\n    if cache is None:\n        cache = {}\n    if term in cache:\n        return cache[term]\n    \n    descendants = set()\n    stack = [term]\n    while stack:\n        cur = stack.pop()\n        for child in go_children.get(cur, []):\n            if child not in descendants:\n                descendants.add(child)\n                stack.append(child)\n    \n    cache[term] = descendants\n    return descendants\n\n\ndef propagate_predictions(\n    predictions_df: pd.DataFrame, \n    go_parents: Dict[str, List[str]]\n) -> pd.DataFrame:\n    \"\"\"\n    Propagate predictions to all ancestor GO terms.\n    Rule: If protein has term X with score S, all ancestors get score >= S (max).\n    \"\"\"\n    print(\"Propagating predictions to ancestor GO terms...\")\n    ancestor_cache = {}\n    propagated_rows = []\n    \n    for pid, group in tqdm(predictions_df.groupby('pid'), desc=\"Propagating\"):\n        term_scores = {}\n        \n        for _, row in group.iterrows():\n            term, score = row['term'], row['p']\n            term_scores[term] = max(term_scores.get(term, 0), score)\n            \n            for ancestor in get_all_ancestors(term, go_parents, ancestor_cache):\n                term_scores[ancestor] = max(term_scores.get(ancestor, 0), score)\n        \n        for term, score in term_scores.items():\n            propagated_rows.append({'pid': pid, 'term': term, 'p': score})\n    \n    propagated_df = pd.DataFrame(propagated_rows)\n    print(f\"Before: {len(predictions_df)}, After: {len(propagated_df)} predictions\")\n    return propagated_df","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_goa_annotations(\n    goa_path: str, \n    go_children: Dict[str, List[str]]\n) -> Tuple[Set[str], Optional[pd.DataFrame]]:\n    \"\"\"\n    Load GOA and build:\n    1. negative_keys: protein-GO pairs to REMOVE (NOT annotations + descendants)\n    2. positive_df: ground truth annotations to ADD with score 1.0\n    \"\"\"\n    if not os.path.exists(goa_path):\n        print(f\"[WARNING] GOA file not found: {goa_path}\")\n        return set(), None\n    \n    print(f\"Loading GOA annotations from {goa_path}...\")\n    goa_df = pd.read_csv(goa_path).drop_duplicates()\n    print(f\"Loaded {len(goa_df)} GOA annotations\")\n    \n    # Extract NEGATIVE annotations\n    print(\"Processing negative annotations...\")\n    neg = goa_df[goa_df['qualifier'].str.contains('NOT', na=False)]\n    neg = neg[['protein_id', 'go_term']].drop_duplicates()\n    neg_by_protein = neg.groupby('protein_id')['go_term'].apply(list).to_dict()\n    \n    # Propagate negatives to descendants\n    desc_cache = {}\n    negative_keys = set()\n    for protein, terms in tqdm(neg_by_protein.items(), desc=\"Propagating negatives\"):\n        all_neg = set(terms)\n        for term in terms:\n            all_neg |= get_all_descendants(term, go_children, desc_cache)\n        for term in all_neg:\n            negative_keys.add(f\"{protein}_{term}\")\n    print(f\"Total negative pairs: {len(negative_keys)}\")\n    \n    # Extract POSITIVE annotations\n    print(\"Processing positive annotations...\")\n    pos = goa_df[~goa_df['qualifier'].str.contains('NOT', na=False)]\n    pos = pos[['protein_id', 'go_term']].drop_duplicates()\n    pos['score'] = 1.0\n    pos['pred_key'] = pos['protein_id'].astype(str) + '_' + pos['go_term'].astype(str)\n    pos = pos[~pos['pred_key'].isin(negative_keys)]\n    print(f\"Total positive pairs: {len(pos)}\")\n    \n    return negative_keys, pos\n\n\ndef apply_negative_propagation(\n    predictions_df: pd.DataFrame, \n    negative_keys: Set[str]\n) -> pd.DataFrame:\n    \"\"\"Remove predictions matching negative protein-GO pairs.\"\"\"\n    if not negative_keys:\n        return predictions_df\n    \n    print(\"Removing negative predictions...\")\n    df = predictions_df.copy()\n    df['pred_key'] = df['pid'].astype(str) + '_' + df['term'].astype(str)\n    before = len(df)\n    df = df[~df['pred_key'].isin(negative_keys)]\n    df = df.drop(columns=['pred_key'])\n    print(f\"Removed {before - len(df)} predictions ({before} -> {len(df)})\")\n    return df\n\n\ndef add_goa_ground_truth(\n    predictions_df: pd.DataFrame, \n    goa_positive_df: Optional[pd.DataFrame]\n) -> pd.DataFrame:\n    \"\"\"Add GOA ground truth annotations with score 1.0.\"\"\"\n    if goa_positive_df is None or len(goa_positive_df) == 0:\n        return predictions_df\n    \n    print(\"Adding GOA ground truth annotations...\")\n    test_proteins = set(predictions_df['pid'].unique())\n    \n    goa = goa_positive_df[goa_positive_df['protein_id'].isin(test_proteins)].copy()\n    goa = goa.rename(columns={'protein_id': 'pid', 'go_term': 'term', 'score': 'p'})\n    goa = goa[['pid', 'term', 'p']]\n    print(f\"Found {len(goa)} GOA annotations for test proteins\")\n    \n    combined = pd.concat([predictions_df, goa], ignore_index=True)\n    combined = combined.groupby(['pid', 'term'])['p'].max().reset_index()\n    print(f\"After adding GOA: {len(predictions_df)} -> {len(combined)}\")\n    return combined","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def merge_goa(file_path = \"/kaggle/input/best-submission/submission1.tsv\"):\n    df = pd.read_csv(file_path, sep=\"\\t\", header=None)\n    \n\n    df = df.rename(columns={0: 'pid', 1: 'term', 2: 'p'})\n    \n    go_parents, go_children = parse_obo(config.obo_path)\n    \n    #load\n    negative_keys, goa_positive_df = load_goa_annotations(config.goa_path, go_children)\n    \n    #propagteted ancestor and descendant\n\n    propagated_df = propagate_predictions(df, go_parents)\n    \n    #negative propaggation # Can stop here for 0.313 model\n    cleaned_df = apply_negative_propagation(propagated_df, negative_keys)\n    final_df = cleaned_df\n    # adding ground truth only work for 0.3 model\n    #final_df = add_goa_ground_truth(cleaned_df, goa_positive_df)\n\n    final_df = final_df.sort_values(by='p', ascending=False)\n    final_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"merge_goa(\"submission.tsv\")\ndf = pd.read_csv(\"submission.tsv\", sep = \"\\t\", header = None)\ndf_sorted = df.sort_values(\n    by=[0, 2], \n    ascending=[True, False]  \n)\ndf_sorted.to_csv(\n    'submission.tsv',  \n    sep='\\t',          \n    header=False,     \n    index=False        \n)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}