{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"},{"sourceId":14187577,"sourceType":"datasetVersion","datasetId":9045748}],"dockerImageVersionId":31240,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"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,"execution":{"iopub.status.busy":"2026-01-10T15:47:00.158149Z","iopub.execute_input":"2026-01-10T15:47:00.158388Z","iopub.status.idle":"2026-01-10T15:47:01.565104Z","shell.execute_reply.started":"2026-01-10T15:47:00.158370Z","shell.execute_reply":"2026-01-10T15:47:01.564369Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- CELL 0: INSTALLATIONS & PATH CHECK ---\n\n# 1. Install Biopython (Crucial for reading FASTA files)\nprint(\"🛠️ Installing Biopython...\")\n!pip install -q biopython\n\n# 2. Check for GPU\nimport torch\nif not torch.cuda.is_available():\n    print(\"❌ WARNING: GPU not detected. Please enable T4 x2 in 'Settings' > 'Accelerator'.\")\nelse:\n    print(f\"✅ GPU Detected: {torch.cuda.get_device_name(0)}\")\n\n# 3. Verify Input Directories (CAFA 6 Dataset)\nimport os\nrequired_files = [\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\",\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n    \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\"\n]\n\nmissing = [f for f in required_files if not os.path.exists(f)]\nif missing:\n    print(f\"❌ CRITICAL ERROR: CAFA 6 Dataset missing! {missing}\")\n    print(\"👉 Please click '+ Add Input' > Search 'cafa-6-protein-function-prediction' > Add.\")\nelse:\n    print(\"✅ CAFA 6 Dataset found.\")\n\n# 4. Verify GOA Dataset (For Negative Propagation)\ngoa_path = \"/kaggle/input/protein-go-annotations/goa_uniprot_all.csv\"\nif os.path.exists(goa_path):\n    print(\"✅ GOA Dataset found (Negative Propagation will work).\")\nelse:\n    print(\"⚠️ GOA Dataset NOT found.\")\n    print(\"👉 ACTION REQUIRED: Click '+ Add Input' -> Search 'sexysoham/protein-go-annotations' -> Add.\")\n    print(\"   (The code will still run without it, but you won't get the Negative Prop score boost.)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T15:47:01.566829Z","iopub.execute_input":"2026-01-10T15:47:01.567109Z","iopub.status.idle":"2026-01-10T15:47:10.437252Z","shell.execute_reply.started":"2026-01-10T15:47:01.567091Z","shell.execute_reply":"2026-01-10T15:47:10.436543Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn\nimport numpy as np\nimport pandas as pd\nfrom transformers import AutoTokenizer, AutoModel\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nimport gc\nfrom Bio import SeqIO\n\n# --- CONFIGURATION ---\nCONFIG = {\n    \"model_name\": \"facebook/esm2_t36_3B_UR50D\", # The 3 Billion Parameter Model\n    \"device\": \"cuda\",\n    \"num_labels\": 2000,       # Top 2000 Terms\n    \"batch_size\": 16,         # Batch size for 3B model (Sorted)\n    \"paths\": {\n        \"train_seq\": \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\",\n        \"train_terms\": \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n        \"test_seq\": \"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\",\n        \"goa_file\": \"/kaggle/input/protein-go-annotations/goa_uniprot_all.csv\" # <--- The new dataset\n    }\n}\n\nprint(f\"🚀 Job Configured: {CONFIG['model_name']} on {torch.cuda.device_count()} GPUs\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T15:47:10.438050Z","iopub.execute_input":"2026-01-10T15:47:10.438421Z","iopub.status.idle":"2026-01-10T15:47:17.002396Z","shell.execute_reply.started":"2026-01-10T15:47:10.438400Z","shell.execute_reply":"2026-01-10T15:47:17.001606Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- STEP 1: FEATURE EXTRACTION (Sorted & Optimized) ---\n\ndef load_fasta(path):\n    seqs = {}\n    for record in SeqIO.parse(path, \"fasta\"):\n        pid = record.id.split(\"|\")[1] if \"|\" in record.id else record.id\n        seqs[pid] = str(record.seq)\n    return seqs\n\ndef extract_embeddings(seq_map, ids, save_name):\n    # Check if exists first\n    if os.path.exists(f\"{save_name}_emb.npy\"):\n        print(f\"✅ Found {save_name}_emb.npy. Loading from disk...\")\n        return np.load(f\"{save_name}_emb.npy\"), np.load(f\"{save_name}_ids.npy\", allow_pickle=True)\n\n    print(f\"🚀 Starting Extraction for {save_name}...\")\n    \n    # SETUP MODEL\n    tokenizer = AutoTokenizer.from_pretrained(CONFIG['model_name'])\n    base_model = AutoModel.from_pretrained(CONFIG['model_name'], torch_dtype=torch.float16)\n    \n    if torch.cuda.device_count() > 1:\n        esm_model = nn.DataParallel(base_model)\n    else:\n        esm_model = base_model\n    esm_model.to(CONFIG['device'])\n    esm_model.eval()\n    \n    # SORT BY LENGTH (Speed Hack)\n    print(\"Sorting sequences by length...\")\n    ids_sorted = sorted(ids, key=lambda x: len(seq_map[x]))\n    \n    embeddings = []\n    \n    with torch.no_grad():\n        for i in tqdm(range(0, len(ids_sorted), CONFIG['batch_size'])):\n            batch_ids = ids_sorted[i : i + CONFIG['batch_size']]\n            batch_seqs = [seq_map[pid] for pid in batch_ids]\n            \n            inputs = tokenizer(\n                batch_seqs, padding=True, truncation=True, max_length=512, return_tensors=\"pt\"\n            )\n            inputs = {k: v.to(CONFIG['device']) for k, v in inputs.items()}\n            \n            outputs = esm_model(**inputs)\n            batch_emb = outputs.last_hidden_state.mean(dim=1).float().cpu().numpy()\n            embeddings.append(batch_emb)\n            \n    full_embeddings = np.concatenate(embeddings, axis=0)\n    \n    # Save\n    np.save(f\"{save_name}_emb.npy\", full_embeddings)\n    np.save(f\"{save_name}_ids.npy\", np.array(ids_sorted))\n    \n    # Cleanup\n    del esm_model, base_model, tokenizer\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return full_embeddings, np.array(ids_sorted)\n\n# RUN EXTRACTION\nprint(\"Loading Data...\")\ntrain_seqs = load_fasta(CONFIG['paths']['train_seq'])\ntest_seqs = load_fasta(CONFIG['paths']['test_seq'])\ntrain_ids = list(train_seqs.keys())\ntest_ids = list(test_seqs.keys())\n\n# Train Data\nX_train, train_ids_sorted = extract_embeddings(train_seqs, train_ids, \"train\")\n\n# Test Data\nX_test, test_ids_sorted = extract_embeddings(test_seqs, test_ids, \"test\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-10T15:47:17.003342Z","iopub.execute_input":"2026-01-10T15:47:17.003904Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- STEP 2: TRAIN CLASSIFIER HEAD ---\nprint(\"🧠 Preparing Training Data...\")\ntrain_terms = pd.read_csv(CONFIG['paths']['train_terms'], sep=\"\\t\")\ntop_terms = train_terms['term'].value_counts().index[:CONFIG['num_labels']].tolist()\nterm_to_idx = {t: i for i, t in enumerate(top_terms)}\nidx_to_term = {i: t for i, t in enumerate(top_terms)}\nprot2terms = train_terms[train_terms['term'].isin(top_terms)].groupby('EntryID')['term'].apply(set).to_dict()\n\nclass EmbeddingDataset(Dataset):\n    def __init__(self, embeddings, pids, labels_dict=None, num_classes=2000):\n        self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n        self.pids = pids\n        self.labels_dict = labels_dict\n        self.num_classes = num_classes\n    def __len__(self): return len(self.pids)\n    def __getitem__(self, idx):\n        emb = self.embeddings[idx]\n        pid = self.pids[idx]\n        label_vec = torch.zeros(self.num_classes)\n        if self.labels_dict and pid in self.labels_dict:\n            for term in self.labels_dict[pid]:\n                if term in term_to_idx: label_vec[term_to_idx[term]] = 1.0\n        return emb, label_vec, pid\n\n# DataLoader\ntrain_ds = EmbeddingDataset(X_train, train_ids_sorted, prot2terms, CONFIG['num_labels'])\ntrain_dl = DataLoader(train_ds, batch_size=256, shuffle=True)\n\n# Model Head\nmodel = nn.Sequential(\n    nn.Linear(2560, 1024), nn.BatchNorm1d(1024), nn.ReLU(), nn.Dropout(0.5),\n    nn.Linear(1024, 512), nn.BatchNorm1d(512), nn.ReLU(), nn.Dropout(0.3),\n    nn.Linear(512, CONFIG['num_labels'])\n).to(CONFIG['device'])\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)\ncriterion = nn.BCEWithLogitsLoss()\n\nprint(\"🚀 Starting Training (20 Epochs)...\")\nfor epoch in range(20):\n    model.train()\n    total_loss = 0\n    for emb, labels, _ in train_dl:\n        emb, labels = emb.to(CONFIG['device']), labels.to(CONFIG['device'])\n        optimizer.zero_grad()\n        loss = criterion(model(emb), labels)\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n    print(f\"Epoch {epoch+1} | Loss: {total_loss/len(train_dl):.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- STEP 3: INFERENCE ---\nprint(\"📝 Generating Predictions...\")\ntest_ds = EmbeddingDataset(X_test, test_ids_sorted, None, CONFIG['num_labels'])\ntest_dl = DataLoader(test_ds, batch_size=256, shuffle=False)\n\nmodel.eval()\npreds = []\nwith torch.no_grad():\n    for emb, _, pids in tqdm(test_dl):\n        emb = emb.to(CONFIG['device'])\n        probs = torch.sigmoid(model(emb)).cpu().numpy()\n        for i, pid in enumerate(pids):\n            # Only keep reasonable scores to save memory\n            indices = np.where(probs[i] > 0.01)[0]\n            for idx in indices:\n                preds.append([pid, idx_to_term[idx], probs[i][idx]])\n\ndf_dl = pd.DataFrame(preds, columns=[\"EntryID\", \"term\", \"score\"])\ndf_dl.to_csv(\"submission_dl_3B.tsv\", sep=\"\\t\", header=False, index=False)\nprint(\"✅ DL Predictions Saved.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- STEP 4: SMART MERGE + NEGATIVE PROPAGATION ---\nprint(\"🚀 Starting Final Optimization Pipeline...\")\n\n# 1. SETUP DIAMOND\nif not os.path.exists(\"diamond\"):\n    !wget -q http://github.com/bbuchfink/diamond/releases/download/v2.1.8/diamond-linux64.tar.gz\n    !tar xzf diamond-linux64.tar.gz\n    !rm diamond-linux64.tar.gz\n    !chmod +x diamond\n    !./diamond makedb --in {CONFIG['paths']['train_seq']} -d train_db --quiet\n\n# 2. RUN DIAMOND (Strict)\nprint(\"💎 Running Diamond (Strict Mode)...\")\n!./diamond blastp -d train_db -q {CONFIG['paths']['test_seq']} -o diamond_matches.tsv -k 3 --min-score 50 --quiet --outfmt 6 qseqid sseqid pident\n\n# Process Diamond\nmatches = pd.read_csv(\"diamond_matches.tsv\", sep=\"\\t\", header=None, names=[\"TestID\", \"TrainID\", \"Identity\"])\ntrain_terms_df = pd.read_csv(CONFIG['paths']['train_terms'], sep=\"\\t\")\ntrain_dict = train_terms_df.groupby(\"EntryID\")[\"term\"].apply(list).to_dict()\n\ndiamond_data = []\nfor _, row in matches.iterrows():\n    tid = row[\"TestID\"].split(\"|\")[1] if \"|\" in str(row[\"TestID\"]) else str(row[\"TestID\"])\n    trid = row[\"TrainID\"].split(\"|\")[1] if \"|\" in str(row[\"TrainID\"]) else str(row[\"TrainID\"])\n    ident = row[\"Identity\"] / 100.0\n    \n    if ident >= 0.6: # Strict Filter\n        if trid in train_dict:\n            for term in train_dict[trid]:\n                diamond_data.append([tid, term, ident])\n\ndf_diamond = pd.DataFrame(diamond_data, columns=[\"EntryID\", \"term\", \"score\"])\nif not df_diamond.empty:\n    df_diamond = df_diamond.groupby([\"EntryID\", \"term\"])[\"score\"].max().reset_index()\n\n# 3. BLEND (DL + Diamond)\nprint(\"⚗️ Blending Models...\")\ncombined = pd.concat([df_dl, df_diamond])\nfinal_df = combined.groupby([\"EntryID\", \"term\"])[\"score\"].max().reset_index()\n\n# 4. NEGATIVE PROPAGATION (The Top-Scorer Trick)\n# Load Negative Annotations (From the dataset you added)\nprint(\"⛔ Applying Negative Propagation...\")\ntry:\n    if os.path.exists(CONFIG['paths']['goa_file']):\n        goa = pd.read_csv(CONFIG['paths']['goa_file'], usecols=['protein_id', 'go_term', 'qualifier'])\n        neg_goa = goa[goa['qualifier'].str.contains(\"NOT\", na=False)]\n        forbidden = set(zip(neg_goa['protein_id'], neg_goa['go_term']))\n        \n        # Filter\n        final_df['key'] = list(zip(final_df['EntryID'], final_df['term']))\n        before_len = len(final_df)\n        final_df = final_df[~final_df['key'].isin(forbidden)]\n        print(f\"   ✂️ Removed {before_len - len(final_df)} impossible terms.\")\n        final_df = final_df.drop(columns=['key'])\n    else:\n        print(\"⚠️ GOA file not found. Skipping Negative Prop.\")\nexcept Exception as e:\n    print(f\"⚠️ Negative Prop Failed: {e}\")\n\n# 5. FINAL FILTERS (Top 75 & Score > 0.15)\nprint(\"✂️ Applying Final Cuts (Top 75 & >0.15)...\")\nfinal_df = final_df.sort_values([\"EntryID\", \"score\"], ascending=[True, False])\nfinal_df = final_df.groupby(\"EntryID\").head(75)\nfinal_df = final_df[final_df[\"score\"] > 0.15]\n\n# Save\nfinal_df.to_csv(\"submission_final.tsv\", sep=\"\\t\", index=False, header=False)\nprint(\"🏆 DONE! 'submission_final.tsv' is ready for submission.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}