{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CAFA-6: Protein Function Prediction","metadata":{}},{"cell_type":"markdown","source":"## Import Libraries","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport os\nimport gc\nfrom collections import Counter, defaultdict\nfrom tqdm.auto import tqdm\n\nfrom Bio import SeqIO\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import StandardScaler, MultiLabelBinarizer\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import TensorDataset, DataLoader\n\nprint(\"[Environment Setup]\")\nprint(f\"PyTorch: {torch.__version__}\")\nprint(f\"CUDA: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"GPUs: {torch.cuda.device_count()}\")\n    for i in range(torch.cuda.device_count()):\n        print(f\"  GPU {i}: {torch.cuda.get_device_name(i)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:41.881372Z","iopub.execute_input":"2026-01-12T10:32:41.881594Z","iopub.status.idle":"2026-01-12T10:32:46.631585Z","shell.execute_reply.started":"2026-01-12T10:32:41.881569Z","shell.execute_reply":"2026-01-12T10:32:46.630946Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Configuration","metadata":{}},{"cell_type":"code","source":"class CFG:\n    BASE_PATH = '/kaggle/input/cafa-6-protein-function-prediction'\n    TRAIN_SEQS = f'{BASE_PATH}/Train/train_sequences.fasta'\n    TRAIN_TERMS = f'{BASE_PATH}/Train/train_terms.tsv'\n    TRAIN_TAX = f'{BASE_PATH}/Train/train_taxonomy.tsv'\n    TEST_SEQS = f'{BASE_PATH}/Test/testsuperset.fasta'\n    TEST_TAX = f'{BASE_PATH}/Test/testsuperset-taxon-list.tsv'\n    OBO_FILE = f'{BASE_PATH}/Train/go-basic.obo'\n    IA_WEIGHTS = f'{BASE_PATH}/IA.tsv'\n    \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    USE_MULTI_GPU = torch.cuda.device_count() > 1\n    \n    FEATURE_DIM = 150\n    TAXONOMY_DIM = 20\n    \n    MLP_EPOCHS = 45\n    MLP_BATCH_SIZE = 512\n    MLP_PATIENCE = 15\n    MLP_LR = 0.002\n    \n    VAL_SIZE = 0.15\n    NUM_MODELS = 2\n    \n    THRESHOLDS = {'F': 0.08, 'P': 0.04, 'C': 0.08}\n    TOP_K = 250\n    ROOTS = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\n\nprint(\"[Configuration]\")\nprint(f\"Device: {CFG.DEVICE}\")\nprint(f\"Multi-GPU: {CFG.USE_MULTI_GPU}\")\nprint(f\"Features: {CFG.FEATURE_DIM + CFG.TAXONOMY_DIM}\")\nprint(f\"Batch size: {CFG.MLP_BATCH_SIZE}\")\nprint(f\"MLP epochs: {CFG.MLP_EPOCHS}\")\nprint(f\"Models per aspect: {CFG.NUM_MODELS}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:46.632508Z","iopub.execute_input":"2026-01-12T10:32:46.632846Z","iopub.status.idle":"2026-01-12T10:32:46.639521Z","shell.execute_reply.started":"2026-01-12T10:32:46.632826Z","shell.execute_reply":"2026-01-12T10:32:46.638775Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load Data","metadata":{}},{"cell_type":"code","source":"print(\"[Data Loading]\")\n\ntrain_terms = pd.read_csv(CFG.TRAIN_TERMS, sep='\\t')\nprint(f\"\\nAnnotations: {len(train_terms):,}\")\n\ntrain_tax = pd.read_csv(CFG.TRAIN_TAX, sep='\\t', header=None, names=['EntryID', 'taxid'])\ntest_tax = pd.read_csv(CFG.TEST_TAX, sep='\\t', header=None, names=['EntryID', 'taxid'])\nprint(f\"Train taxonomy: {len(train_tax):,}\")\nprint(f\"Test taxonomy: {len(test_tax):,}\")\n\nia_weights = pd.read_csv(CFG.IA_WEIGHTS, sep='\\t')\nia_dict = dict(zip(ia_weights.iloc[:, 0], ia_weights.iloc[:, 1]))\nprint(f\"IA weights: {len(ia_dict):,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:46.640437Z","iopub.execute_input":"2026-01-12T10:32:46.640722Z","iopub.status.idle":"2026-01-12T10:32:47.061870Z","shell.execute_reply.started":"2026-01-12T10:32:46.640701Z","shell.execute_reply":"2026-01-12T10:32:47.061231Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Parse FASTA Files","metadata":{}},{"cell_type":"code","source":"def parse_fasta(filepath):\n    sequences = {}\n    for record in SeqIO.parse(filepath, 'fasta'):\n        pid = record.id.split('|')[1] if '|' in record.id else record.id.split()[0]\n        sequences[pid] = str(record.seq)\n    return sequences\n\nprint(\"[Parsing FASTA]\")\n\ntrain_seqs = parse_fasta(CFG.TRAIN_SEQS)\nprint(f\"\\nTrain: {len(train_seqs):,}\")\n\ntest_seqs = parse_fasta(CFG.TEST_SEQS)\nprint(f\"Test: {len(test_seqs):,}\")\n\ntrain_lens = [len(seq) for seq in train_seqs.values()]\ntest_lens = [len(seq) for seq in test_seqs.values()]\n\nprint(f\"\\nTrain length: {np.mean(train_lens):.0f} +/- {np.std(train_lens):.0f}\")\nprint(f\"Test length: {np.mean(test_lens):.0f} +/- {np.std(test_lens):.0f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:47.063828Z","iopub.execute_input":"2026-01-12T10:32:47.064099Z","iopub.status.idle":"2026-01-12T10:32:51.282417Z","shell.execute_reply.started":"2026-01-12T10:32:47.064077Z","shell.execute_reply":"2026-01-12T10:32:51.281717Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Exploratory Data Analysis","metadata":{}},{"cell_type":"code","source":"print(\"[Exploratory Data Analysis]\")\n\naspect_counts = train_terms['aspect'].value_counts()\naspect_names = {'F': 'Molecular Function', 'P': 'Biological Process', 'C': 'Cellular Component'}\n\nprint(\"\\nAnnotations per aspect:\")\nfor asp in ['F', 'P', 'C']:\n    print(f\"{aspect_names[asp]}: {aspect_counts[asp]:,}\")\n\nprint(\"\\nUnique terms per aspect:\")\nfor asp in ['F', 'P', 'C']:\n    n = train_terms[train_terms['aspect'] == asp]['term'].nunique()\n    print(f\"{aspect_names[asp]}: {n:,}\")\n\nproteins = train_terms['EntryID'].nunique()\nprint(f\"\\nProteins: {proteins:,}\")\nprint(f\"Avg terms per protein: {len(train_terms) / proteins:.1f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:51.283350Z","iopub.execute_input":"2026-01-12T10:32:51.283654Z","iopub.status.idle":"2026-01-12T10:32:51.529645Z","shell.execute_reply.started":"2026-01-12T10:32:51.283624Z","shell.execute_reply":"2026-01-12T10:32:51.529079Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, axes = plt.subplots(2, 2, figsize=(14, 10))\n\naxes[0, 0].hist(train_lens, bins=50, alpha=0.7, label='Train', color='steelblue', edgecolor='black')\naxes[0, 0].hist(test_lens, bins=50, alpha=0.7, label='Test', color='coral', edgecolor='black')\naxes[0, 0].set_xlabel('Sequence Length')\naxes[0, 0].set_ylabel('Frequency')\naxes[0, 0].set_title('Sequence Length Distribution', fontweight='bold')\naxes[0, 0].legend()\naxes[0, 0].grid(True, alpha=0.3)\n\naspects = [aspect_names[k] for k in ['F', 'P', 'C']]\ncounts = [aspect_counts[k] for k in ['F', 'P', 'C']]\naxes[0, 1].barh(aspects, counts, color=['steelblue', 'coral', 'seagreen'])\naxes[0, 1].set_xlabel('Count')\naxes[0, 1].set_title('Annotations per Aspect', fontweight='bold')\naxes[0, 1].grid(True, alpha=0.3, axis='x')\n\ntop_terms = train_terms['term'].value_counts().head(20)\naxes[1, 0].barh(range(len(top_terms)), top_terms.values, color='steelblue')\naxes[1, 0].set_yticks(range(len(top_terms)))\naxes[1, 0].set_yticklabels(top_terms.index, fontsize=8)\naxes[1, 0].set_xlabel('Count')\naxes[1, 0].set_title('Top 20 GO Terms', fontweight='bold')\naxes[1, 0].invert_yaxis()\naxes[1, 0].grid(True, alpha=0.3, axis='x')\n\ntop_tax = train_tax['taxid'].value_counts().head(15)\naxes[1, 1].barh(range(len(top_tax)), top_tax.values, color='coral')\naxes[1, 1].set_yticks(range(len(top_tax)))\naxes[1, 1].set_yticklabels(top_tax.index, fontsize=8)\naxes[1, 1].set_xlabel('Count')\naxes[1, 1].set_title('Top 15 Species', fontweight='bold')\naxes[1, 1].invert_yaxis()\naxes[1, 1].grid(True, alpha=0.3, axis='x')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:51.530470Z","iopub.execute_input":"2026-01-12T10:32:51.530759Z","iopub.status.idle":"2026-01-12T10:32:53.054553Z","shell.execute_reply.started":"2026-01-12T10:32:51.530728Z","shell.execute_reply":"2026-01-12T10:32:53.053936Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## GO Ontology Processing","metadata":{}},{"cell_type":"code","source":"def parse_obo(obo_file):\n    term_parents = defaultdict(set)\n    current_term = None\n    \n    with open(obo_file, 'r') as f:\n        for line in f:\n            line = line.strip()\n            if line == '[Term]':\n                current_term = None\n            elif line.startswith('id: '):\n                current_term = line.split('id: ')[1]\n            elif line.startswith('is_a: ') and current_term:\n                parent = line.split('is_a: ')[1].split(' !')[0].strip()\n                term_parents[current_term].add(parent)\n            elif line.startswith('relationship: part_of ') and current_term:\n                parent = line.split('part_of ')[1].split(' !')[0].strip()\n                term_parents[current_term].add(parent)\n    \n    return term_parents\n\ndef get_ancestors_map(term_parents):\n    ancestors = {}\n    \n    def get_all_ancestors(term):\n        if term in ancestors:\n            return ancestors[term]\n        parents = term_parents.get(term, set())\n        all_anc = set(parents)\n        for parent in parents:\n            all_anc |= get_all_ancestors(parent)\n        ancestors[term] = all_anc\n        return all_anc\n    \n    for term in tqdm(list(term_parents.keys()), desc='Building ancestors'):\n        get_all_ancestors(term)\n    \n    return ancestors\n\nprint(\"[GO Ontology Processing]\")\n\nterm_parents = parse_obo(CFG.OBO_FILE)\nprint(f\"\\nTerms: {len(term_parents):,}\")\n\nancestors_map = get_ancestors_map(term_parents)\nprint(f\"Ancestors: {len(ancestors_map):,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:53.055413Z","iopub.execute_input":"2026-01-12T10:32:53.055708Z","iopub.status.idle":"2026-01-12T10:32:53.658728Z","shell.execute_reply.started":"2026-01-12T10:32:53.055682Z","shell.execute_reply":"2026-01-12T10:32:53.658045Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Enhanced Feature Engineering","metadata":{}},{"cell_type":"code","source":"class EnhancedFeatureExtractor:\n    @staticmethod\n    def extract(seq):\n        if not seq:\n            return np.zeros(150, dtype=np.float32)\n        \n        try:\n            length = len(seq)\n            aa_counts = Counter(seq)\n            \n            standard = 'ACDEFGHIKLMNPQRSTVWY'\n            aa_freq = np.array([aa_counts.get(aa, 0)/length for aa in standard], dtype=np.float32)\n            \n            hydro = sum(aa_counts.get(aa, 0) for aa in 'AILMFWYV') / length\n            polar = sum(aa_counts.get(aa, 0) for aa in 'STNQ') / length\n            charged = sum(aa_counts.get(aa, 0) for aa in 'DEKR') / length\n            pos = sum(aa_counts.get(aa, 0) for aa in 'RK') / length\n            neg = sum(aa_counts.get(aa, 0) for aa in 'DE') / length\n            arom = sum(aa_counts.get(aa, 0) for aa in 'FWY') / length\n            aliph = sum(aa_counts.get(aa, 0) for aa in 'ILV') / length\n            small = sum(aa_counts.get(aa, 0) for aa in 'ACDGNPSTV') / length\n            tiny = sum(aa_counts.get(aa, 0) for aa in 'ACGS') / length\n            \n            biochem = np.array([np.log1p(length), hydro, polar, charged, pos, neg,\n                               arom, aliph, small, tiny], dtype=np.float32)\n            \n            top_di = ['AL','LA','LE','EA','AA','AS','SA','EL','LL','AE','SE','ES','GA','AG','VA','AV','LV','VL','LS','SL',\n                     'RE','ER','VE','EV','TE','ET','GL','LG','KE','EK','AR','RA','KA','AK','GE','EG','LK','KL','IA','AI']\n            di_freq = np.zeros(40, dtype=np.float32)\n            if length > 1:\n                di_counts = Counter([seq[i:i+2] for i in range(length-1)])\n                di_freq = np.array([di_counts.get(d, 0)/(length-1) for d in top_di], dtype=np.float32)\n            \n            top_tri = ['ALA','LEA','EAL','LAL','AAA','LLE','ELE','ALE','GAL','ASA','VLA','LAV','SLS','LSL','GLA','LAG',\n                      'AVL','VLA','SLE','LES','AAL','LAA','AAS','SAA','GGA','AGG','GGG','AAE','EAA','LAE',\n                      'KEA','AEL','AKE','KEL','ELA','ALK','LKE','KAL','VAL','ALV']\n            tri_freq = np.zeros(40, dtype=np.float32)\n            if length > 2:\n                tri_counts = Counter([seq[i:i+3] for i in range(length-2)])\n                tri_freq = np.array([tri_counts.get(t, 0)/(length-2) for t in top_tri], dtype=np.float32)\n            \n            n_term = seq[:min(50, length)]\n            c_term = seq[-min(50, length):]\n            \n            n_hydro = sum(1 for aa in n_term if aa in 'AILMFWYV') / len(n_term)\n            c_hydro = sum(1 for aa in c_term if aa in 'AILMFWYV') / len(c_term)\n            n_charged = sum(1 for aa in n_term if aa in 'DEKR') / len(n_term)\n            c_charged = sum(1 for aa in c_term if aa in 'DEKR') / len(c_term)\n            \n            unique = len(set(seq)) / 20\n            low_comp = max(aa_counts.values()) / length\n            helix = sum(aa_counts.get(aa, 0) for aa in 'AELM') / length\n            sheet = sum(aa_counts.get(aa, 0) for aa in 'VIY') / length\n            turn = sum(aa_counts.get(aa, 0) for aa in 'NPGS') / length\n            net_charge = pos - neg\n            charge_dens = charged / np.log1p(length)\n            hydro_mom = hydro * np.log1p(length)\n            flex = sum(aa_counts.get(aa, 0) for aa in 'GQSPE') / length\n            rigid = sum(aa_counts.get(aa, 0) for aa in 'WYCFIL') / length\n            disorder = sum(aa_counts.get(aa, 0) for aa in 'AEQKPS') / length\n            instab = (hydro + charged) / 2\n            iso = (pos * 10 + neg * 4) / (pos + neg + 1)\n            \n            positional = np.array([n_hydro, c_hydro, n_charged, c_charged, unique, low_comp,\n                                  helix, sheet, turn, net_charge, charge_dens, hydro_mom,\n                                  flex, rigid, disorder, instab, iso, 0, 0, 0], dtype=np.float32)\n            \n            feats = np.concatenate([aa_freq, biochem, di_freq, tri_freq, positional])\n            return np.pad(feats, (0, max(0, 150-len(feats))), 'constant')[:150]\n        \n        except:\n            return np.zeros(150, dtype=np.float32)\n\nprint(\"[Feature Extraction]\")\n\ntrain_features = {}\nfor pid, seq in tqdm(train_seqs.items(), desc='Train features'):\n    train_features[pid] = EnhancedFeatureExtractor.extract(seq)\n\ntest_features = {}\nfor pid, seq in tqdm(test_seqs.items(), desc='Test features'):\n    test_features[pid] = EnhancedFeatureExtractor.extract(seq)\n\nprint(f\"\\nTrain: {len(train_features):,}\")\nprint(f\"Test: {len(test_features):,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:32:53.659778Z","iopub.execute_input":"2026-01-12T10:32:53.660336Z","iopub.status.idle":"2026-01-12T10:34:29.122553Z","shell.execute_reply.started":"2026-01-12T10:32:53.660305Z","shell.execute_reply":"2026-01-12T10:34:29.121727Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Taxonomy Encoding & Scaling","metadata":{}},{"cell_type":"code","source":"print(\"[Taxonomy Encoding & Scaling]\")\n\ntop_taxa = train_tax['taxid'].value_counts().head(CFG.TAXONOMY_DIM).index.tolist()\n\ndef get_tax_vector(pid, tax_map, top_taxa):\n    vec = np.zeros(len(top_taxa), dtype=np.float32)\n    if pid in tax_map and tax_map[pid] in top_taxa:\n        vec[top_taxa.index(tax_map[pid])] = 1.0\n    return vec\n\ntrain_tax_map = train_tax.set_index('EntryID')['taxid'].to_dict()\ntest_tax_map = test_tax.set_index('EntryID')['taxid'].to_dict()\n\ntrain_pids = list(train_features.keys())\ntest_pids = list(test_features.keys())\n\nX_train = np.array([train_features[p] for p in train_pids])\nX_test = np.array([test_features[p] for p in test_pids])\n\nscaler = StandardScaler()\nX_train = scaler.fit_transform(X_train)\nX_test = scaler.transform(X_test)\n\nfinal_train_features = {}\nfor i, pid in enumerate(train_pids):\n    final_train_features[pid] = np.concatenate([X_train[i], get_tax_vector(pid, train_tax_map, top_taxa)])\n\nfinal_test_features = {}\nfor i, pid in enumerate(test_pids):\n    final_test_features[pid] = np.concatenate([X_test[i], get_tax_vector(pid, test_tax_map, top_taxa)])\n\nprint(f\"\\nFinal dimension: {CFG.FEATURE_DIM + CFG.TAXONOMY_DIM}\")\n\ndel train_features, test_features, X_train, X_test\ngc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:34:29.123494Z","iopub.execute_input":"2026-01-12T10:34:29.123708Z","iopub.status.idle":"2026-01-12T10:34:30.841355Z","shell.execute_reply.started":"2026-01-12T10:34:29.123688Z","shell.execute_reply":"2026-01-12T10:34:30.840814Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset Preparation","metadata":{}},{"cell_type":"code","source":"print(\"[Dataset Preparation]\")\n\naspect_data = {}\naspect_names = {'F': 'Molecular Function', 'P': 'Biological Process', 'C': 'Cellular Component'}\n\nfor code, name in aspect_names.items():\n    print(f\"\\n[{name}]\")\n    \n    aspect_terms = train_terms[train_terms['aspect'] == code]\n    prot_terms = aspect_terms.groupby('EntryID')['term'].apply(list).to_dict()\n    valid_pids = [p for p in prot_terms.keys() if p in final_train_features]\n    \n    mlb = MultiLabelBinarizer(sparse_output=False)\n    y = mlb.fit_transform([prot_terms[p] for p in valid_pids])\n    X = np.array([final_train_features[p] for p in valid_pids], dtype=np.float32)\n    \n    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=CFG.VAL_SIZE, random_state=42)\n    \n    print(f\"Terms: {len(mlb.classes_):,}\")\n    print(f\"Train: {X_train.shape[0]:,}\")\n    print(f\"Val: {X_val.shape[0]:,}\")\n    \n    aspect_data[code] = {'X_train': X_train, 'X_val': X_val, 'y_train': y_train, 'y_val': y_val, 'mlb': mlb}\n\nprint(\"\\nDataset preparation complete\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:34:30.842130Z","iopub.execute_input":"2026-01-12T10:34:30.842371Z","iopub.status.idle":"2026-01-12T10:34:41.619656Z","shell.execute_reply.started":"2026-01-12T10:34:30.842345Z","shell.execute_reply":"2026-01-12T10:34:41.618992Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Definitions","metadata":{}},{"cell_type":"code","source":"class MLP(nn.Module):\n    def __init__(self, input_dim, num_classes, hidden_dim=256, dropout=0.4):\n        super(MLP, self).__init__()\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, hidden_dim//2),\n            nn.BatchNorm1d(hidden_dim//2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim//2, num_classes)\n        )\n    \n    def forward(self, x):\n        return self.net(x)\n\ndef train_mlp(model, train_loader, val_loader, epochs, device, use_multi_gpu, model_name):\n    if use_multi_gpu and torch.cuda.device_count() > 1:\n        model = nn.DataParallel(model)\n        print(f\"Using {torch.cuda.device_count()} GPUs\")\n    \n    model = model.to(device)\n    criterion = nn.BCEWithLogitsLoss()\n    optimizer = torch.optim.AdamW(model.parameters(), lr=CFG.MLP_LR, weight_decay=1e-4)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n    \n    best_loss = float('inf')\n    patience = 0\n    best_state = None\n    \n    pbar = tqdm(range(epochs), desc=f'{model_name} training')\n    for epoch in pbar:\n        model.train()\n        train_loss = 0\n        for X_batch, y_batch in train_loader:\n            X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n            optimizer.zero_grad()\n            loss = criterion(model(X_batch), y_batch)\n            loss.backward()\n            optimizer.step()\n            train_loss += loss.item()\n        \n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for X_batch, y_batch in val_loader:\n                X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n                val_loss += criterion(model(X_batch), y_batch).item()\n        \n        train_loss /= len(train_loader)\n        val_loss /= len(val_loader)\n        scheduler.step()\n        \n        if val_loss < best_loss:\n            best_loss = val_loss\n            patience = 0\n            best_state = model.module.state_dict().copy() if use_multi_gpu else model.state_dict().copy()\n        else:\n            patience += 1\n        \n        pbar.set_postfix({'train': f'{train_loss:.4f}', 'val': f'{val_loss:.4f}', 'best': f'{best_loss:.4f}'})\n        \n        if patience >= CFG.MLP_PATIENCE:\n            print(f\" Early stop at epoch {epoch+1}\")\n            break\n    \n    if best_state:\n        if use_multi_gpu and isinstance(model, nn.DataParallel):\n            model.module.load_state_dict(best_state)\n        else:\n            model.load_state_dict(best_state)\n    \n    return model.module if use_multi_gpu and isinstance(model, nn.DataParallel) else model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:34:41.620548Z","iopub.execute_input":"2026-01-12T10:34:41.620813Z","iopub.status.idle":"2026-01-12T10:34:41.633447Z","shell.execute_reply.started":"2026-01-12T10:34:41.620778Z","shell.execute_reply":"2026-01-12T10:34:41.632852Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Training","metadata":{}},{"cell_type":"code","source":"all_models = {}\n\nfor code, name in aspect_names.items():\n    print(f\"[Training {name}]\")\n    \n    data = aspect_data[code]\n    X_train, X_val = data['X_train'], data['X_val']\n    y_train, y_val = data['y_train'], data['y_val']\n    num_classes = y_train.shape[1]\n    \n    models = []\n    \n    train_ds = TensorDataset(torch.FloatTensor(X_train), torch.FloatTensor(y_train))\n    val_ds = TensorDataset(torch.FloatTensor(X_val), torch.FloatTensor(y_val))\n    train_loader = DataLoader(train_ds, batch_size=CFG.MLP_BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True)\n    val_loader = DataLoader(val_ds, batch_size=CFG.MLP_BATCH_SIZE, num_workers=2, pin_memory=True)\n    \n    for i in range(CFG.NUM_MODELS):\n        print(f\"\\n[Model {i+1}/{CFG.NUM_MODELS}]\")\n        hidden_dim = [256, 320][i % 2]\n        mlp = MLP(X_train.shape[1], num_classes, hidden_dim=hidden_dim)\n        mlp = train_mlp(mlp, train_loader, val_loader, CFG.MLP_EPOCHS, CFG.DEVICE, CFG.USE_MULTI_GPU, f\"MLP-{i+1}\")\n        models.append(mlp)\n    \n    del train_ds, val_ds, train_loader, val_loader\n    gc.collect()\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n    \n    all_models[code] = models\n    print(f\"\\nCompleted {name}\")\n\nprint(\"Training complete\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:34:41.634223Z","iopub.execute_input":"2026-01-12T10:34:41.634457Z","iopub.status.idle":"2026-01-12T10:45:24.415764Z","shell.execute_reply.started":"2026-01-12T10:34:41.634439Z","shell.execute_reply":"2026-01-12T10:45:24.414945Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Test Predictions","metadata":{}},{"cell_type":"code","source":"print(\"[Test Predictions & Post-Processing]\")\n\ntest_pids_ordered = [p for p in test_pids if p in final_test_features]\nprint(f\"\\nTest proteins: {len(test_pids_ordered):,}\")\n\nfinal_predictions = []\n\nPROTEIN_CHUNK_SIZE = 15000\n\ndef process_predictions_inline(pid, scores_dict, ancestors_map, roots):\n    updated = scores_dict.copy()\n    for term, score in scores_dict.items():\n        if term in ancestors_map:\n            for anc in ancestors_map[term]:\n                updated[anc] = max(updated.get(anc, 0.0), score)\n    non_root = [s for t, s in updated.items() if t not in roots]\n    if non_root:\n        max_val = max(non_root)\n        if 0 < max_val < 0.95:\n            for t in updated:\n                if t not in roots:\n                    updated[t] = min(1.0, updated[t] * (0.95 / max_val))\n    if updated:\n        for r in roots:\n            updated[r] = 1.0\n    return updated\n\nfor code, name in aspect_names.items():\n    print(f\"\\n[{name}]\")\n    \n    models = all_models[code]\n    mlb = aspect_data[code]['mlb']\n    num_classes = len(mlb.classes_)\n    thresh = CFG.THRESHOLDS[code]\n    \n    num_protein_chunks = (len(test_pids_ordered) + PROTEIN_CHUNK_SIZE - 1) // PROTEIN_CHUNK_SIZE\n    \n    for chunk_idx in range(num_protein_chunks):\n        start_protein = chunk_idx * PROTEIN_CHUNK_SIZE\n        end_protein = min(start_protein + PROTEIN_CHUNK_SIZE, len(test_pids_ordered))\n        chunk_pids = test_pids_ordered[start_protein:end_protein]\n        \n        X_chunk = np.array([final_test_features[p] for p in chunk_pids], dtype=np.float32)\n        \n        print(f\"  Chunk {chunk_idx+1}/{num_protein_chunks}: {len(chunk_pids):,} proteins\")\n        \n        chunk_preds = np.zeros((len(chunk_pids), num_classes), dtype=np.float32)\n        \n        for model_idx, model in enumerate(models):\n            model.eval()\n            model = model.to(CFG.DEVICE)\n            \n            batch_size = 2048\n            num_batches = (len(X_chunk) + batch_size - 1) // batch_size\n            \n            with torch.no_grad():\n                for batch_idx in range(num_batches):\n                    start_idx = batch_idx * batch_size\n                    end_idx = min(start_idx + batch_size, len(X_chunk))\n                    \n                    batch_data = torch.FloatTensor(X_chunk[start_idx:end_idx]).to(CFG.DEVICE)\n                    batch_out = torch.sigmoid(model(batch_data)).cpu().numpy()\n                    chunk_preds[start_idx:end_idx] += batch_out\n                    \n                    del batch_data, batch_out\n            \n            model.cpu()\n            gc.collect()\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n        \n        chunk_preds /= len(models)\n        \n        for i, pid in enumerate(chunk_pids):\n            scores = chunk_preds[i]\n            \n            term_scores = {}\n            for j, s in enumerate(scores):\n                if s >= thresh:\n                    term_scores[mlb.classes_[j]] = float(s)\n            \n            if len(term_scores) < 15:\n                top_idx = np.argsort(scores)[-30:]\n                for j in top_idx:\n                    term_scores[mlb.classes_[j]] = float(scores[j])\n            \n            processed = process_predictions_inline(pid, term_scores, ancestors_map, CFG.ROOTS)\n            sorted_terms = sorted(processed.items(), key=lambda x: -x[1])[:CFG.TOP_K]\n            \n            for term, score in sorted_terms:\n                if score >= 0.001:\n                    final_predictions.append((pid, term, score))\n        \n        del X_chunk, chunk_preds\n        gc.collect()\n\ngc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()\n\nprint(f\"\\nTotal predictions: {len(final_predictions):,}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T10:45:24.418400Z","iopub.execute_input":"2026-01-12T10:45:24.418756Z","iopub.status.idle":"2026-01-12T11:11:18.106069Z","shell.execute_reply.started":"2026-01-12T10:45:24.418731Z","shell.execute_reply":"2026-01-12T11:11:18.105512Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Create Submission","metadata":{}},{"cell_type":"code","source":"print(\"[Creating Submission]\")\n\ndf = pd.DataFrame(final_predictions, columns=['protein_id', 'go_term', 'score'])\ndf = df.sort_values(['protein_id', 'score'], ascending=[True, False])\n\nprint(f\"\\nTotal predictions: {len(df):,}\")\nprint(f\"Unique proteins: {df['protein_id'].nunique():,}\")\nprint(f\"Unique terms: {df['go_term'].nunique():,}\")\n\nprint(f\"\\nScore range: [{df['score'].min():.6f}, {df['score'].max():.6f}]\")\n\ndups = df.duplicated(subset=['protein_id', 'go_term']).sum()\nif dups:\n    df = df.drop_duplicates(subset=['protein_id', 'go_term'])\n    print(f\"Removed {dups} duplicates\")\n\nmax_pp = df.groupby('protein_id').size().max()\nprint(f\"Max per protein: {max_pp}\")\n\ndf.to_csv('submission.tsv', sep='\\t', index=False, header=False)\nsize = os.path.getsize('submission.tsv') / (1024*1024)\nprint(f\"\\nFile: submission.tsv ({size:.1f} MB)\")\n\nprint(\"\\nFirst 10 predictions:\")\nprint(df.head(10).to_string(index=False))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T11:11:18.111541Z","iopub.execute_input":"2026-01-12T11:11:18.111980Z","iopub.status.idle":"2026-01-12T11:15:09.195653Z","shell.execute_reply.started":"2026-01-12T11:11:18.111960Z","shell.execute_reply":"2026-01-12T11:15:09.195005Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Statistics","metadata":{}},{"cell_type":"code","source":"print(\"[Statistics]\")\n\nprint(f\"\\nScore mean: {df['score'].mean():.4f}\")\nprint(f\"Score std: {df['score'].std():.4f}\")\nprint(f\"Score median: {df['score'].median():.4f}\")\n\nppp = df.groupby('protein_id').size()\nprint(f\"\\nPer protein mean: {ppp.mean():.1f}\")\nprint(f\"Per protein std: {ppp.std():.1f}\")\nprint(f\"Per protein range: [{ppp.min()}, {ppp.max()}]\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T11:15:09.197001Z","iopub.execute_input":"2026-01-12T11:15:09.197602Z","iopub.status.idle":"2026-01-12T11:15:15.916096Z","shell.execute_reply.started":"2026-01-12T11:15:09.197579Z","shell.execute_reply":"2026-01-12T11:15:15.915455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\naxes[0].hist(df['score'], bins=50, color='steelblue', edgecolor='black')\naxes[0].set_xlabel('Score')\naxes[0].set_ylabel('Frequency')\naxes[0].set_title('Score Distribution', fontweight='bold')\naxes[0].grid(True, alpha=0.3)\n\naxes[1].hist(ppp, bins=50, color='coral', edgecolor='black')\naxes[1].set_xlabel('Predictions per Protein')\naxes[1].set_ylabel('Frequency')\naxes[1].set_title('Predictions per Protein', fontweight='bold')\naxes[1].grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T11:15:15.916887Z","iopub.execute_input":"2026-01-12T11:15:15.917187Z","iopub.status.idle":"2026-01-12T11:15:17.133279Z","shell.execute_reply.started":"2026-01-12T11:15:15.917159Z","shell.execute_reply":"2026-01-12T11:15:17.132667Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Summary","metadata":{}},{"cell_type":"code","source":"print(\"[Summary]\")\n\nprint(f\"\\nTrain proteins: {len(train_seqs):,}\")\nprint(f\"Test proteins: {len(test_seqs):,}\")\nprint(f\"GO terms: {len(term_parents):,}\")\n\nprint(f\"\\nFeatures: {CFG.FEATURE_DIM + CFG.TAXONOMY_DIM}\")\nprint(f\"Models: {len(all_models) * CFG.NUM_MODELS}\")\nprint(f\"GPUs used: {torch.cuda.device_count() if torch.cuda.is_available() else 0}\")\n\nprint(f\"\\nPredictions: {len(df):,}\")\nprint(f\"File: submission.tsv ({size:.1f} MB)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-12T11:15:17.134168Z","iopub.execute_input":"2026-01-12T11:15:17.134447Z","iopub.status.idle":"2026-01-12T11:15:17.139893Z","shell.execute_reply.started":"2026-01-12T11:15:17.134418Z","shell.execute_reply":"2026-01-12T11:15:17.139089Z"}},"outputs":[],"execution_count":null}]}