{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ==========================================\n# 1. SETUP & LIBRARIES\n# ==========================================\n!pip install -q obonet biopython transformers torch numpy pandas tqdm scikit-learn\n\nimport os\nimport sys\nimport torch\nimport torch.nn as nn\nimport numpy as np\nimport pandas as pd\nimport obonet\nimport networkx as nx\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\nfrom transformers import AutoTokenizer, AutoModel\nfrom Bio import SeqIO\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.auto import tqdm\nimport gc\n\n# --- CONFIGURATION ---\nCONFIG = {\n    # Paths\n    'seq_file': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta',\n    'term_file': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv',\n    'tax_file': '/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv',\n    'obo_file': '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo',\n    'test_fasta': '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta',\n    'test_tax': '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv',\n    'save_path': 'model_checkpoint.pth',\n    'data_cache': 'processed_data.pt',\n    \n    # Model Parameters\n    'model_name': 'facebook/esm2_t6_8M_UR50D', # Switch to 'facebook/esm2_t12_35M_UR50D' if you have >12GB VRAM\n    'num_labels': 1500,    # Increased from 1000 for better coverage\n    'max_len': 512,        # Max sequence length\n    \n    # Training Hyperparameters\n    'batch_size': 32,      # Increased thanks to Mixed Precision\n    'lr': 5e-4,\n    'epochs': 5,\n    'device': 'cuda' if torch.cuda.is_available() else 'cpu',\n    'num_workers': 2       # Parallel data loading\n}\n\nprint(f\"🚀 Hardware: {CONFIG['device']}\")\nif CONFIG['device'] == 'cuda':\n    print(f\"   GPU: {torch.cuda.get_device_name(0)}\")\n\n# ==========================================\n# 2. OPTIMIZED DATA PROCESSING\n# ==========================================\n\ndef load_and_process_data():\n    \"\"\"Loads raw data, propagates labels, and caches the result.\"\"\"\n    \n    # Check if cache exists to save time\n    if os.path.exists(CONFIG['data_cache']):\n        print(f\"\\n[Cache] Loading processed data from {CONFIG['data_cache']}...\")\n        return torch.load(CONFIG['data_cache'])\n\n    print(\"\\n[1/4] Reading GO Graph & Taxonomy...\")\n    go_graph = obonet.read_obo(CONFIG['obo_file'])\n    \n    # Helper: Taxonomy Map\n    tax_df = pd.read_csv(CONFIG['tax_file'], sep='\\t', header=None, names=['id', 'tax_id'])\n    tax_df['tax_id'] = tax_df['tax_id'].astype(str)\n    unique_taxa = sorted(tax_df['tax_id'].unique().tolist()) # Sorted for consistency\n    tax_to_idx = {t: i for i, t in enumerate(unique_taxa)}\n    \n    # Helper: Term Map (Top N)\n    print(f\"[2/4] Selecting Top {CONFIG['num_labels']} Frequent Terms...\")\n    terms_df = pd.read_csv(CONFIG['term_file'], sep='\\t', header=None, names=['id', 'term', 'aspect'])\n    top_terms = terms_df['term'].value_counts().head(CONFIG['num_labels']).index.tolist()\n    term_to_idx = {t: i for i, t in enumerate(top_terms)}\n    \n    # Pre-calculate ancestors for the top terms only (optimization)\n    # We only care about ancestors if they are ALSO in our top_terms list\n    allowed_terms = set(top_terms)\n    \n    print(\"[3/4] Processing Sequences & Propagating Labels...\")\n    data_dict = {}\n    \n    # 1. Load Seqs\n    for record in SeqIO.parse(CONFIG['seq_file'], \"fasta\"):\n        pid = record.id.split('|')[1] if '|' in record.id else record.id\n        data_dict[pid] = {'seq': str(record.seq), 'labels': set(), 'tax': 0}\n        \n    # 2. Map Tax\n    for _, row in tax_df.iterrows():\n        if row['id'] in data_dict:\n            data_dict[row['id']]['tax'] = tax_to_idx.get(row['tax_id'], 0)\n            \n    # 3. Map Labels (Fast Propagation)\n    grouped_terms = terms_df.groupby('id')['term'].apply(list)\n    \n    for pid, terms in tqdm(grouped_terms.items(), desc=\"Propagating GO terms\"):\n        if pid in data_dict:\n            propagated = set()\n            for t in terms:\n                if t in go_graph:\n                    # Get ancestors, but immediately filter by allowed_terms to save memory\n                    ancestors = nx.descendants(go_graph, t) | {t}\n                    propagated.update(ancestors & allowed_terms)\n            \n            # Convert to indices\n            data_dict[pid]['labels'] = [term_to_idx[t] for t in propagated]\n\n    # 4. Final List Structure\n    final_data = []\n    for pid, info in data_dict.items():\n        if info['labels']: # Only keep if it has labels\n            final_data.append({\n                'seq': info['seq'],\n                'tax': info['tax'],\n                'labels': info['labels'], # Store as list of indices (sparse), not dense vector yet\n                'id': pid\n            })\n            \n    meta_data = {\n        'data': final_data,\n        'tax_to_idx': tax_to_idx,\n        'term_to_idx': term_to_idx,\n        'idx_to_term': {i: t for t, i in term_to_idx.items()}\n    }\n    \n    # Save cache\n    print(f\"-> Saving processed data ({len(final_data)} samples) to cache...\")\n    torch.save(meta_data, CONFIG['data_cache'])\n    return meta_data\n\n# Execute Data Loading\ndb = load_and_process_data()\nfull_data = db['data']\ntax_to_idx = db['tax_to_idx']\nidx_to_term = db['idx_to_term']\nnum_taxa = len(tax_to_idx)\n\n# ==========================================\n# 3. DATASET & MODEL (Memory Efficient)\n# ==========================================\n\nclass ProteinDataset(Dataset):\n    def __init__(self, data_list, tokenizer, num_labels, max_len=512):\n        self.data = data_list\n        self.tokenizer = tokenizer\n        self.max_len = max_len\n        self.num_labels = num_labels\n\n    def __len__(self): return len(self.data)\n\n    def __getitem__(self, idx):\n        item = self.data[idx]\n        \n        # Dynamic Tokenization (Faster than pre-tokenizing everything into RAM)\n        inputs = self.tokenizer(\n            item['seq'], \n            return_tensors=\"pt\", \n            padding='max_length', \n            truncation=True, \n            max_length=self.max_len\n        )\n        \n        # Create Dense Label Vector on the fly (Saves RAM)\n        label_vec = torch.zeros(self.num_labels, dtype=torch.float32)\n        label_vec[item['labels']] = 1.0\n        \n        return {\n            'ids': inputs['input_ids'].squeeze(0),\n            'mask': inputs['attention_mask'].squeeze(0),\n            'tax': torch.tensor(item['tax'], dtype=torch.long),\n            'labels': label_vec\n        }\n\nclass EfficientProteinPredictor(nn.Module):\n    def __init__(self, model_name, num_classes, num_taxa, tax_embed_dim=64):\n        super().__init__()\n        self.esm = AutoModel.from_pretrained(model_name)\n        \n        # Freezing ESM is critical for speed and preventing overfitting on small data\n        for param in self.esm.parameters():\n            param.requires_grad = False\n            \n        self.esm_dim = self.esm.config.hidden_size\n        self.tax_embedding = nn.Embedding(num_taxa, tax_embed_dim)\n        \n        # Lightweight Head\n        self.classifier = nn.Sequential(\n            nn.Linear(self.esm_dim + tax_embed_dim, 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, input_ids, attention_mask, tax_ids):\n        # ESM Forward (Feature Extraction)\n        with torch.no_grad(): # Ensure no gradients for ESM body\n            esm_out = self.esm(input_ids=input_ids, attention_mask=attention_mask)\n        \n        # Get CLS token\n        seq_emb = esm_out.last_hidden_state[:, 0, :] \n        tax_emb = self.tax_embedding(tax_ids)\n        \n        combined = torch.cat((seq_emb, tax_emb), dim=1)\n        return self.classifier(combined)\n\n# ==========================================\n# 4. TRAINING WITH MIXED PRECISION (AMP)\n# ==========================================\n\n# Setup\ntrain_data, val_data = train_test_split(full_data, test_size=0.1, random_state=42)\ntokenizer = AutoTokenizer.from_pretrained(CONFIG['model_name'])\n\ntrain_ds = ProteinDataset(train_data, tokenizer, CONFIG['num_labels'], CONFIG['max_len'])\nval_ds = ProteinDataset(val_data, tokenizer, CONFIG['num_labels'], CONFIG['max_len'])\n\ntrain_loader = DataLoader(train_ds, batch_size=CONFIG['batch_size'], shuffle=True, num_workers=CONFIG['num_workers'], pin_memory=True)\nval_loader = DataLoader(val_ds, batch_size=CONFIG['batch_size'], num_workers=CONFIG['num_workers'], pin_memory=True)\n\nmodel = EfficientProteinPredictor(CONFIG['model_name'], CONFIG['num_labels'], num_taxa).to(CONFIG['device'])\noptimizer = torch.optim.AdamW(model.parameters(), lr=CONFIG['lr'], weight_decay=0.01)\ncriterion = nn.BCEWithLogitsLoss()\nscaler = GradScaler() # For Mixed Precision\n\nprint(f\"\\n[Training] Batches per epoch: {len(train_loader)}\")\n\nbest_val_loss = float('inf')\n\nfor epoch in range(CONFIG['epochs']):\n    model.train()\n    total_loss = 0\n    \n    # Progress bar\n    pbar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{CONFIG['epochs']}\")\n    \n    for batch in pbar:\n        ids = batch['ids'].to(CONFIG['device'], non_blocking=True)\n        mask = batch['mask'].to(CONFIG['device'], non_blocking=True)\n        tax = batch['tax'].to(CONFIG['device'], non_blocking=True)\n        labels = batch['labels'].to(CONFIG['device'], non_blocking=True)\n        \n        optimizer.zero_grad()\n        \n        # --- Mixed Precision Block ---\n        with autocast(): \n            logits = model(ids, mask, tax)\n            loss = criterion(logits, labels)\n            \n        scaler.scale(loss).backward()\n        scaler.step(optimizer)\n        scaler.update()\n        # -----------------------------\n        \n        total_loss += loss.item()\n        pbar.set_postfix({'loss': loss.item()})\n        \n    # Validation\n    model.eval()\n    val_loss = 0\n    with torch.no_grad():\n        for batch in val_loader:\n            ids = batch['ids'].to(CONFIG['device'])\n            mask = batch['mask'].to(CONFIG['device'])\n            tax = batch['tax'].to(CONFIG['device'])\n            labels = batch['labels'].to(CONFIG['device'])\n            \n            with autocast(): # Faster inference\n                logits = model(ids, mask, tax)\n                loss = criterion(logits, labels)\n            val_loss += loss.item()\n            \n    avg_val_loss = val_loss / len(val_loader)\n    print(f\"Epoch {epoch+1} | Train Loss: {total_loss/len(train_loader):.4f} | Val Loss: {avg_val_loss:.4f}\")\n    \n    # Save Best\n    if avg_val_loss < best_val_loss:\n        best_val_loss = avg_val_loss\n        torch.save(model.state_dict(), CONFIG['save_path'])\n        print(\"-> Model Saved!\")\n\n# ==========================================\n# 5. FAST VECTORIZED INFERENCE\n# ==========================================\n\nprint(\"\\n[Inference] Loading Test Data...\")\n\n# Reuse the training taxonomy logic for test set\ntest_tax_df = pd.read_csv(CONFIG['test_tax'], sep='\\t', header=None, names=['id', 'tax_id'])\ntest_tax_map = {row['id']: tax_to_idx.get(str(row['tax_id']), 0) for _, row in test_tax_df.iterrows()}\n\nclass SubmissionDataset(Dataset):\n    def __init__(self, fasta_file, tax_map, tokenizer, max_len=512):\n        self.data = []\n        self.tokenizer = tokenizer\n        self.max_len = max_len\n        \n        # Fast Parsing\n        for record in SeqIO.parse(fasta_file, \"fasta\"):\n            pid = record.id.split('|')[1] if '|' in record.id else record.id\n            self.data.append({\n                'id': pid,\n                'seq': str(record.seq),\n                'tax': tax_map.get(pid, 0)\n            })\n\n    def __len__(self): return len(self.data)\n    \n    def __getitem__(self, idx):\n        item = self.data[idx]\n        inputs = self.tokenizer(item['seq'], return_tensors=\"pt\", padding='max_length', truncation=True, max_length=self.max_len)\n        return {\n            'ids': inputs['input_ids'].squeeze(0),\n            'mask': inputs['attention_mask'].squeeze(0),\n            'tax': torch.tensor(item['tax'], dtype=torch.long),\n            'pid': item['id']\n        }\n\nsub_ds = SubmissionDataset(CONFIG['test_fasta'], test_tax_map, tokenizer)\nsub_loader = DataLoader(sub_ds, batch_size=CONFIG['batch_size']*2, shuffle=False, num_workers=2) # Double batch size for inference\n\n# Load Best Model\nmodel.load_state_dict(torch.load(CONFIG['save_path']))\nmodel.eval()\n\nsubmission_file = \"submission.tsv\"\nTOP_K = 150 # Keeping this realistic (Rules allow 1500, but storing 1500 per protein creates huge files. 150 is usually enough for top score)\nMIN_SCORE = 0.01\n\nprint(f\"Generating predictions for {len(sub_ds)} proteins...\")\n\nwith open(submission_file, 'w') as f:\n    with torch.no_grad():\n        for batch in tqdm(sub_loader, desc=\"Predicting\"):\n            ids = batch['ids'].to(CONFIG['device'])\n            mask = batch['mask'].to(CONFIG['device'])\n            tax = batch['tax'].to(CONFIG['device'])\n            pids = batch['pid']\n            \n            with autocast():\n                logits = model(ids, mask, tax)\n                probs = torch.sigmoid(logits)\n            \n            # --- VECTORIZED TOP-K (HUGE SPEEDUP) ---\n            # Instead of loop-sort, we get top K values and indices directly on GPU\n            top_vals, top_inds = torch.topk(probs, k=TOP_K, dim=1)\n            \n            # Move to CPU for writing\n            top_vals = top_vals.cpu().numpy()\n            top_inds = top_inds.cpu().numpy()\n            \n            # Write Buffer\n            buffer = []\n            for i in range(len(pids)):\n                pid = pids[i]\n                for j in range(TOP_K):\n                    score = top_vals[i][j]\n                    if score < MIN_SCORE: break # Since it's sorted, we can break early\n                    \n                    term = idx_to_term[top_inds[i][j]]\n                    buffer.append(f\"{pid}\\t{term}\\t{score:.3f}\\n\")\n            \n            # Bulk Write\n            f.writelines(buffer)\n\nprint(f\"\\n✅ Done! Submission saved to {submission_file}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T11:30:37.152510Z","iopub.execute_input":"2026-01-02T11:30:37.153078Z","iopub.status.idle":"2026-01-02T12:05:35.455133Z","shell.execute_reply.started":"2026-01-02T11:30:37.153051Z","shell.execute_reply":"2026-01-02T12:05:35.454283Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dfd = pd.read_csv('/kaggle/working/submission.tsv')\ndfd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T12:06:01.520494Z","iopub.execute_input":"2026-01-02T12:06:01.520817Z","iopub.status.idle":"2026-01-02T12:06:10.090116Z","shell.execute_reply.started":"2026-01-02T12:06:01.520793Z","shell.execute_reply":"2026-01-02T12:06:10.089505Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dfd.iloc[1]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T12:06:18.340141Z","iopub.execute_input":"2026-01-02T12:06:18.340675Z","iopub.status.idle":"2026-01-02T12:06:18.346308Z","shell.execute_reply.started":"2026-01-02T12:06:18.340645Z","shell.execute_reply":"2026-01-02T12:06:18.345737Z"}},"outputs":[],"execution_count":null}]}