{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13680623,"sourceType":"datasetVersion","datasetId":8699749},{"sourceId":13966329,"sourceType":"datasetVersion","datasetId":8835088},{"sourceId":683492,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":518643,"modelId":533185}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nfrom torch.utils.data import DataLoader, random_split, Dataset\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport pandas as pd\nimport numpy as np\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom tqdm import tqdm\nfrom torchmetrics.classification import MultilabelF1Score\nimport matplotlib.pyplot as plt\nfrom collections import defaultdict\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n","metadata":{"_uuid":"1acdacfd-307f-4f6a-afbf-d2143a0a68fe","_cell_guid":"7cd7aaaa-f772-4354-933a-9650b79d9a90","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-12-13T04:55:28.061097Z","iopub.execute_input":"2025-12-13T04:55:28.061773Z","iopub.status.idle":"2025-12-13T04:55:39.576367Z","shell.execute_reply.started":"2025-12-13T04:55:28.061738Z","shell.execute_reply":"2025-12-13T04:55:39.575529Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Path configurations\n","metadata":{}},{"cell_type":"code","source":"IS_LOCAL = not os.path.isdir(\"/kaggle/input\")\nCHECKPOINT_PATH = \"/kaggle/input/checkpoint-310-cafa6/pytorch/default/1\"\nTRAIN_TERMS_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv'\nTRAIN_SEQUENCES_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\nTEST_SEQUENCES_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\nPROT_EMBEDS = '/kaggle/input/esm-650m-embeds/esm-650m-embeds/protein_embeddings.npy'\nPIDS = '/kaggle/input/esm-650m-embeds/esm-650m-embeds/protein_ids.csv'\nOBO_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo'\nGOA_PATH = '/kaggle/input/protein-go-annotations/goa_uniprot_all.csv'\n\nif IS_LOCAL:\n    CHECKPOINT_PATH = \"./models\"\n    TRAIN_TERMS_PATH = './data/cafa-6-protein-function-prediction/Train/train_terms.tsv'\n    TRAIN_SEQUENCES_PATH = './data/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\n    TEST_SEQUENCES_PATH = './data/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\n    PIDS = \"./data/cafa-6-protein-function-prediction/esm-650m-embeds/protein_ids.csv\"\n    PROT_EMBEDS = \"./data/cafa-6-protein-function-prediction/esm-650m-embeds/protein_embeddings.npy\"\n    OBO_PATH = './data/cafa-6-protein-function-prediction/Train/go-basic.obo'\n    GOA_PATH = './data/cafa-6-protein-function-prediction/goa_uniprot_all.csv' \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:39.577701Z","iopub.execute_input":"2025-12-13T04:55:39.578062Z","iopub.status.idle":"2025-12-13T04:55:39.582989Z","shell.execute_reply.started":"2025-12-13T04:55:39.578044Z","shell.execute_reply":"2025-12-13T04:55:39.582285Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load precalculated embeddings\n","metadata":{}},{"cell_type":"code","source":"protein_ids = pd.read_csv(PIDS)[\"protein_id\"].tolist()\nembeddings = np.load(PROT_EMBEDS)\nembeddings_dict = {pid: emb for pid, emb in zip(protein_ids, embeddings)}\nprint(f\"Loaded {len(protein_ids)} embeddings of dimension {embeddings.shape[1]}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:39.583758Z","iopub.execute_input":"2025-12-13T04:55:39.584568Z","iopub.status.idle":"2025-12-13T04:55:49.714158Z","shell.execute_reply.started":"2025-12-13T04:55:39.584532Z","shell.execute_reply":"2025-12-13T04:55:49.713324Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load data","metadata":{}},{"cell_type":"code","source":"\ndef parse_fasta(fasta_file) -> dict[str, str]:\n    sequences = {}\n    current_id = None\n    current_seq = []\n    with open(fasta_file, 'r') as f:\n        for line in f:\n            line = line.strip()\n            if line.startswith('>'):\n                if current_id:\n                    sequences[current_id] = ''.join(current_seq)\n                parts = line[1:].split('|')\n                if len(parts) >= 2:\n                    current_id = parts[1]\n                else:\n                    current_id = line[1:].split()[0]\n                current_seq = []\n            else:\n                current_seq.append(line)\n        if current_id:\n            sequences[current_id] = ''.join(current_seq)\n    return sequences\n\ntrain_terms_df = pd.read_csv(TRAIN_TERMS_PATH, sep='\\t')\ntrain_sequences = parse_fasta(TRAIN_SEQUENCES_PATH)\ntest_sequences = parse_fasta(TEST_SEQUENCES_PATH)\nprint(f\"Train size: {len(train_sequences)}, Test size: {len(test_sequences)}\")\nprint(f\"Total annotations: {len(train_terms_df)}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:49.715996Z","iopub.execute_input":"2025-12-13T04:55:49.716567Z","iopub.status.idle":"2025-12-13T04:55:52.171235Z","shell.execute_reply.started":"2025-12-13T04:55:49.716546Z","shell.execute_reply":"2025-12-13T04:55:52.170357Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Go Hierarchy Parsing","metadata":{}},{"cell_type":"code","source":"# Thanks to https://www.kaggle.com/code/seddiktrk/cafa-6-blend-goa-negative-propagation\ndef parse_obo(go_obo_path):\n    parents = defaultdict(set)\n    children = defaultdict(set)\n    \n    if not os.path.exists(go_obo_path): \n        return parents, children\n        \n    with open(go_obo_path,\"r\") as f:\n        cur_id=None\n        for line in f:\n            line=line.strip()\n            if line==\"[Term]\": \n                cur_id=None\n            elif line.startswith(\"id: \"): \n                cur_id=line.split(\"id: \")[1].strip()\n            elif line.startswith(\"is_a: \"):\n                pid=line.split()[1].strip()\n                if cur_id: \n                    parents[cur_id].add(pid)\n                    children[pid].add(cur_id)\n            elif line.startswith(\"relationship: part_of \"):\n                parts=line.split(); \n                if len(parts)>=3:\n                    pid=parts[2].strip()\n                    if cur_id: \n                        parents[cur_id].add(pid)\n                        children[pid].add(cur_id)\n    print(f\"[io] Parsed OBO: {len(parents)} nodes with parents\")\n    return parents, children\n\ndef get_all_ancestors(term, go_parents, cache=None):\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\ndef get_all_descendants(term, go_children, cache=None):\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# Load GO hierarchy\ngo_parents, go_children = parse_obo(OBO_PATH)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:52.172152Z","iopub.execute_input":"2025-12-13T04:55:52.172830Z","iopub.status.idle":"2025-12-13T04:55:52.670008Z","shell.execute_reply.started":"2025-12-13T04:55:52.172798Z","shell.execute_reply":"2025-12-13T04:55:52.669349Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Post-processing utility","metadata":{}},{"cell_type":"code","source":"# ===== GOA NEGATIVE PROPAGATION =====\ndef load_goa_and_build_negative_keys(goa_path, go_children):\n    \"\"\"\n    Load GOA annotations and build:\n    1. negative_keys: protein-GO pairs that should be removed (not annotations + their descendants)\n    2. goa_ground_truth: known positive annotations to add with score 1.0\n    \"\"\"\n    goa_df = pd.read_csv(goa_path)\n    goa_df = goa_df.drop_duplicates()\n    print(f\"Loaded {len(goa_df)} GOA annotations\")\n    \n    # 1. Extract NEGATIVE annotations (NOT qualifiers)\n    negative_annots = goa_df[goa_df['qualifier'].str.contains('NOT', na=False)]\n    negative_annots = negative_annots[['protein_id', 'go_term']].drop_duplicates()\n    negative_by_protein = negative_annots.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(negative_by_protein.items(), desc=\"Propagating negatives\"):\n        all_negative_terms = set(terms)\n        for term in terms:\n            all_negative_terms |= get_all_descendants(term, go_children, desc_cache)\n        for term in all_negative_terms:\n            negative_keys.add(f\"{protein}_{term}\")\n    \n    print(f\"Total unique negative protein-GO pairs: {len(negative_keys)}\")\n    \n    # 2. Extract POSITIVE annotations (ground truth)\n    positive_annots = goa_df[~goa_df['qualifier'].str.contains('NOT', na=False)]\n    positive_annots = positive_annots[['protein_id', 'go_term']].drop_duplicates()\n    positive_annots['score'] = 1.0\n    positive_annots['pred_key'] = positive_annots['protein_id'].astype(str) + '_' + positive_annots['go_term'].astype(str)\n    # Remove any that are also negative\n    positive_annots = positive_annots[~positive_annots['pred_key'].isin(negative_keys)]\n    \n    print(f\"Total positive GOA ground truth pairs: {len(positive_annots)}\")\n    \n    return negative_keys, positive_annots\n\n\ndef propagate_predictions(predictions_df, go_parents):\n    print(\"Propagating predictions to ancestor GO terms...\")\n    \n    # Build ancestor cache for efficiency\n    ancestor_cache = {}\n    \n    # Group by protein\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 = row['term']\n            score = row['p']\n            \n            # Add original prediction\n            if term not in term_scores or score > term_scores[term]:\n                term_scores[term] = score\n            \n            # Propagate to all ancestors\n            ancestors = get_all_ancestors(term, go_parents, ancestor_cache)\n            for ancestor in ancestors:\n                if ancestor not in term_scores or score > term_scores[ancestor]:\n                    term_scores[ancestor] = score\n        \n        # Convert back to rows\n        for term, score in term_scores.items():\n            propagated_rows.append({\n                'pid': pid,\n                'term': term,\n                'p': score\n            })\n    \n    propagated_df = pd.DataFrame(propagated_rows)\n    print(f\"Before propagation: {len(predictions_df)} predictions\")\n    print(f\"After propagation: {len(propagated_df)} predictions\")\n    \n    return propagated_df\n\ndef apply_negative_propagation(predictions_df, negative_keys):\n    if not negative_keys:\n        return predictions_df\n    \n    predictions_df['pred_key'] = predictions_df['pid'].astype(str) + '_' + predictions_df['term'].astype(str)\n    before_count = len(predictions_df)\n    predictions_df = predictions_df[~predictions_df['pred_key'].isin(negative_keys)]\n    after_count = len(predictions_df)\n    predictions_df = predictions_df.drop(columns=['pred_key'])\n    \n    print(f\"Removed {before_count - after_count} negative predictions ({before_count} -> {after_count})\")\n    return predictions_df\n\ndef add_goa_ground_truth(predictions_df, goa_positive_df):\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    # Get test proteins from predictions\n    test_proteins = set(predictions_df['pid'].unique())\n    \n    # Filter GOA to only include test proteins\n    goa_for_test = goa_positive_df[goa_positive_df['protein_id'].isin(test_proteins)].copy()\n    goa_for_test = goa_for_test.rename(columns={'protein_id': 'pid', 'go_term': 'term', 'score': 'p'})\n    goa_for_test = goa_for_test[['pid', 'term', 'p']]\n    \n    print(f\"Found {len(goa_for_test)} GOA annotations for test proteins\")\n    \n    # Combine, keeping max score for duplicates\n    combined = pd.concat([predictions_df, goa_for_test], ignore_index=True)\n    combined = combined.groupby(['pid', 'term'])['p'].max().reset_index()\n    \n    print(f\"After adding GOA: {len(predictions_df)} -> {len(combined)} predictions\")\n    return combined","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:52.670718Z","iopub.execute_input":"2025-12-13T04:55:52.670946Z","iopub.status.idle":"2025-12-13T04:55:52.909016Z","shell.execute_reply.started":"2025-12-13T04:55:52.670928Z","shell.execute_reply":"2025-12-13T04:55:52.908343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"negative_keys, goa_positive_df = load_goa_and_build_negative_keys(GOA_PATH, go_children)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:52.909711Z","iopub.execute_input":"2025-12-13T04:55:52.910484Z","iopub.status.idle":"2025-12-13T04:55:58.658861Z","shell.execute_reply.started":"2025-12-13T04:55:52.910465Z","shell.execute_reply":"2025-12-13T04:55:58.658122Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"#  Dataset Classes","metadata":{}},{"cell_type":"code","source":"class ProtDataset(Dataset):\n    def __init__(self, pids, labels, embeddings_dict):\n        self.pids = pids\n        self.labels = labels\n        self.embeddings_dict = embeddings_dict\n        \n    def __len__(self):\n        return len(self.pids)\n\n    def __getitem__(self, idx):\n        embed = self.embeddings_dict[self.pids[idx]]\n        if hasattr(self.labels, 'toarray'):\n            label_array = self.labels[idx].toarray().flatten()\n        else:\n            label_array = self.labels[idx]\n        label_tensor = torch.tensor(label_array, dtype=torch.float32)\n        return torch.from_numpy(embed).float(), label_tensor\n\nclass TestDataSet(Dataset):\n    def __init__(self, pids, embeddings_dict):\n        self.pids = pids\n        self.embeddings_dict = embeddings_dict\n        \n    def __len__(self):\n        return len(self.pids)\n\n    def __getitem__(self, idx):\n        return torch.from_numpy(self.embeddings_dict[self.pids[idx]]).float()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:58.659497Z","iopub.execute_input":"2025-12-13T04:55:58.659706Z","iopub.status.idle":"2025-12-13T04:55:58.666027Z","shell.execute_reply.started":"2025-12-13T04:55:58.659690Z","shell.execute_reply":"2025-12-13T04:55:58.665322Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n# Model Definition\n","metadata":{}},{"cell_type":"code","source":"class SimpleMLP(nn.Module):\n    def __init__(self, input_dim=1280, num_classes=3000):\n        super().__init__()\n        self.network = nn.Sequential(\n            nn.Linear(input_dim, 1024),\n            nn.BatchNorm1d(1024),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(1024, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(512, num_classes)\n        )\n        \n    def forward(self, x):\n        return self.network(x)\n\ndef get_improved_model(num_classes):\n    return SimpleMLP(input_dim=1280, num_classes=num_classes)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:58.666656Z","iopub.execute_input":"2025-12-13T04:55:58.666922Z","iopub.status.idle":"2025-12-13T04:55:58.685771Z","shell.execute_reply.started":"2025-12-13T04:55:58.666898Z","shell.execute_reply":"2025-12-13T04:55:58.685020Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train utility","metadata":{}},{"cell_type":"code","source":"# ===== VISUALIZATION =====\ndef plot_losses_and_scores(train_losses, train_scores, val_scores, aspect_name):\n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n    \n    ax1.plot(train_losses, 'r-', marker='o', label='Train Loss')\n    ax1.set_xlabel('Epoch')\n    ax1.set_ylabel('Loss')\n    ax1.set_title(f'{aspect_name} - Training Loss')\n    ax1.legend()\n    ax1.grid(True)\n    \n    ax2.plot(train_scores, 'b-', marker='o', label='Train F1')\n    ax2.plot(val_scores, 'g-', marker='s', label='Val F1')\n    ax2.set_xlabel('Epoch')\n    ax2.set_ylabel('F1 Score')\n    ax2.set_title(f'{aspect_name} - F1 Scores')\n    ax2.legend()\n    ax2.grid(True)\n    \n    plt.tight_layout()\n    plt.savefig(f'{aspect_name}_training_curves.png', dpi=150)\n    plt.show()\n\ndef train(aspect_name, model, train_loader, valid_loader, num_epochs, num_classes, lr):\n    model = model.to(device)\n    loss_fn = nn.BCEWithLogitsLoss()\n    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n        optimizer, mode='min', factor=0.5, patience=2\n    )\n    \n    torch_f1_score = MultilabelF1Score(num_labels=num_classes, threshold=0.05, average='micro').to(device=device)\n    \n    train_losses = []\n    train_scores = []\n    val_scores = []\n    \n    for epoch in range(1, num_epochs + 1):\n        model.train()\n        train_loss = 0.0\n        \n        for X, y in train_loader:\n            X = X.to(device)\n            y = y.to(device)\n            optimizer.zero_grad()\n            \n            outputs = model(X)\n            loss = loss_fn(outputs, y)\n        \n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item() * X.size(0)\n        \n        avg_train_loss = train_loss / len(train_loader.dataset)\n        train_losses.append(avg_train_loss)\n        \n        # Validation\n        model.eval()\n        all_val_preds = []\n        all_val_labels = []\n        val_loss = 0.0\n        \n        for valid_seqs, valid_labels in valid_loader:\n            valid_seqs = valid_seqs.to(device)\n            valid_labels = valid_labels.to(device)\n            \n            with torch.no_grad():\n                outputs = model(valid_seqs)\n                batch_loss = loss_fn(outputs, valid_labels)\n                val_loss += batch_loss.item() * valid_seqs.size(0)\n                preds = torch.sigmoid(outputs)\n                all_val_preds.append(preds)\n                all_val_labels.append(valid_labels)\n        \n        all_val_preds = torch.vstack(all_val_preds)\n        all_val_labels = torch.vstack(all_val_labels)\n        val_f1 = torch_f1_score(all_val_preds, all_val_labels).item()\n        val_scores.append(val_f1)\n        \n        avg_val_loss = val_loss / len(valid_loader.dataset)\n        scheduler.step(avg_val_loss)\n        \n        # print(f\"Epoch {epoch}/{num_epochs} | Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f} | Val F1: {val_f1:.4f}\")\n    \n    plot_losses_and_scores(train_losses, train_scores, val_scores, aspect_name)\n    \n    optimizer = None\n    scheduler = None\n    torch.cuda.empty_cache()\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:58.687959Z","iopub.execute_input":"2025-12-13T04:55:58.688488Z","iopub.status.idle":"2025-12-13T04:55:58.701752Z","shell.execute_reply.started":"2025-12-13T04:55:58.688462Z","shell.execute_reply":"2025-12-13T04:55:58.701113Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train 3 models for 3 ontologies","metadata":{}},{"cell_type":"code","source":"ASPECTS = ['C', 'F', 'P']\nASPECT_NAMES = {'C': 'Cellular Component (CCO)', 'F': 'Molecular Function (MFO)', 'P': 'Biological Process (BPO)'}\naspect_models = {}\naspect_mlbs = {}\n\ndef train_aspect_model(aspect_name):\n    global aspect_models, aspect_mlbs\n    \n    aspect_df = train_terms_df[train_terms_df['aspect'] == aspect_name]\n    print(f\"Training model for {aspect_name}\")\n    \n    # Group terms by protein for this aspect\n    protein_2_terms = aspect_df.groupby('EntryID')['term'].apply(list).to_dict()\n    \n    # Only include proteins that have:\n    # 1. Embeddings available\n    # 2. At least one GO term for this aspect\n    pid_train = [pid for pid in train_sequences.keys() \n                 if pid in embeddings_dict and pid in protein_2_terms]\n    \n    labels_list = [protein_2_terms[pid] for pid in pid_train]\n    \n    mlb = MultiLabelBinarizer(sparse_output=True)\n    y_train_labels = mlb.fit_transform(labels_list)\n    aspect_mlbs[aspect_name] = mlb\n    \n    num_classes = len(mlb.classes_)\n    print(f\"{aspect_name} - Number of proteins: {y_train_labels.shape[0]}, GO terms: {num_classes}\")\n    \n\n    train_dataset = ProtDataset(pid_train, y_train_labels, embeddings_dict)\n    train_size = int(0.9 * len(train_dataset))\n    valid_size = len(train_dataset) - train_size\n    \n    train_part, valid_part = random_split(train_dataset, [train_size, valid_size])\n    train_loader = DataLoader(train_part, batch_size=128, shuffle=True, num_workers=4)\n    valid_loader = DataLoader(valid_part, batch_size=128, shuffle=False, num_workers=4)\n    \n    num_classes = len(mlb.classes_)\n    epochs = 30\n    lr = 9e-4\n    \n    model = train(\n        aspect_name,\n        get_improved_model(num_classes),\n        train_loader, valid_loader, epochs,\n        num_classes, lr\n    )\n    \n    aspect_models[aspect_name] = model\n    \n    del train_loader, valid_loader, train_dataset, train_part, valid_part\n    del y_train_labels\n    torch.cuda.empty_cache()\n    if torch.cuda.is_available():\n        torch.cuda.synchronize()\n    \n    return model, mlb\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:58.702342Z","iopub.execute_input":"2025-12-13T04:55:58.702561Z","iopub.status.idle":"2025-12-13T04:55:58.720276Z","shell.execute_reply.started":"2025-12-13T04:55:58.702546Z","shell.execute_reply":"2025-12-13T04:55:58.719521Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for aspect in ASPECTS:\n    train_aspect_model(aspect)\n    torch.save({\n        \"model\": aspect_models[aspect].state_dict(),\n    }, f\"checkpoint-{aspect}.pth\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T04:55:58.721051Z","iopub.execute_input":"2025-12-13T04:55:58.721436Z","iopub.status.idle":"2025-12-13T04:56:03.763900Z","shell.execute_reply.started":"2025-12-13T04:55:58.721395Z","shell.execute_reply":"2025-12-13T04:56:03.762810Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Prediction","metadata":{}},{"cell_type":"code","source":"\ndef predict(threshold=0.02, min_predictions=25):\n    \n    pid_test = [pid for pid in test_sequences.keys() if pid in embeddings_dict]\n    \n    test_dataset = TestDataSet(pid_test, embeddings_dict)\n    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=3)\n    \n    submission_list = []\n    \n    for aspect in ASPECTS:\n        model = aspect_models[aspect]\n        mlb = aspect_mlbs[aspect]\n        model.to(device)\n        model.eval()\n        \n        batch_start_idx = 0\n        \n        for test_embeds in tqdm(test_loader, desc=f\"Predicting ({aspect})\"):\n            test_embeds = test_embeds.to(device)\n            batch_size = test_embeds.shape[0]\n            \n            with torch.no_grad():\n                outputs = model(test_embeds)\n                preds = torch.sigmoid(outputs).cpu()\n            \n            for i in range(batch_size):\n                protein_id = pid_test[batch_start_idx + i]\n                pred_scores = preds[i].numpy()\n                \n                pred_term_ids = np.where(pred_scores >= threshold)[0]\n                \n                if len(pred_term_ids) < min_predictions:\n                    pred_term_ids = np.argsort(pred_scores)[-min_predictions:][::-1]\n                \n                for term_idx in pred_term_ids:\n                    submission_list.append({\n                        'pid': protein_id,\n                        'term': mlb.classes_[term_idx],\n                        'p': pred_scores[term_idx].item()\n                    })\n            \n            batch_start_idx += batch_size\n            \n            del test_embeds, outputs, preds\n            if batch_start_idx % 5000 == 0:\n                torch.cuda.empty_cache()\n    \n    submission_df = pd.DataFrame(submission_list)\n    print(f\"\\nTotal predictions: {len(submission_df)}\")\n    \n    return submission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T05:00:56.122358Z","iopub.execute_input":"2025-12-13T05:00:56.122690Z","iopub.status.idle":"2025-12-13T05:00:56.136946Z","shell.execute_reply.started":"2025-12-13T05:00:56.122638Z","shell.execute_reply":"2025-12-13T05:00:56.136198Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_aspect_cls(aspect_name):\n    aspect_df = train_terms_df[train_terms_df['aspect'] == aspect_name]\n    # Group terms by protein for this aspect\n    protein_2_terms = aspect_df.groupby('EntryID')['term'].apply(list).to_dict()\n    \n    # Only include proteins that have:\n    # 1. Embeddings available\n    # 2. At least one GO term for this aspect\n    pid_train = [pid for pid in train_sequences.keys() \n                 if pid in embeddings_dict and pid in protein_2_terms]\n    \n    labels_list = [protein_2_terms[pid] for pid in pid_train]\n    \n    mlb = MultiLabelBinarizer(sparse_output=True)\n    y_train_labels = mlb.fit_transform(labels_list)\n    aspect_mlbs[aspect_name] = mlb\n    \n    return len(mlb.classes_)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T05:00:27.668568Z","iopub.execute_input":"2025-12-13T05:00:27.669402Z","iopub.status.idle":"2025-12-13T05:00:27.674247Z","shell.execute_reply.started":"2025-12-13T05:00:27.669374Z","shell.execute_reply":"2025-12-13T05:00:27.673291Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# for aspect in ASPECTS:\n#     checkpoint = torch.load(f\"{CHECKPOINT_PATH}/checkpoint-{aspect}.pth\", map_location=\"cpu\")\n#     num_classes = get_aspect_cls(aspect)\n#     aspect_models[aspect]= get_improved_model(num_classes)\n#     aspect_models[aspect].load_state_dict(checkpoint[\"model\"])\n\ncombined_submission_df = predict(threshold=0.02)\npropagated_df = propagate_predictions(combined_submission_df, go_parents)\ncleaned_df = apply_negative_propagation(propagated_df, negative_keys)\nfinal_df = add_goa_ground_truth(cleaned_df, goa_positive_df)\n# final_df = final_df.sort_values(by='p', ascending=False)\nfinal_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-13T05:00:58.333081Z","iopub.execute_input":"2025-12-13T05:00:58.333784Z","iopub.status.idle":"2025-12-13T05:23:25.845522Z","shell.execute_reply.started":"2025-12-13T05:00:58.333749Z","shell.execute_reply":"2025-12-13T05:23:25.843619Z"}},"outputs":[],"execution_count":null}]}