{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31260,"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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install biopython obonet --quiet","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:01:25.870505Z","iopub.execute_input":"2026-01-24T12:01:25.87083Z","iopub.status.idle":"2026-01-24T12:01:30.449536Z","shell.execute_reply.started":"2026-01-24T12:01:25.870796Z","shell.execute_reply":"2026-01-24T12:01:30.448822Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 1. Install Missing Libraries (if needed) ---\n!pip install biopython obonet --quiet\n\n# --- 2. Imports ---\nimport obonet\nimport networkx as nx\nimport os\nfrom Bio import SeqIO\n\n# --- 3. Define File Paths (Updated based on your screenshot) ---\n# Root directory\nDATA_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\n\n# Note: The files are inside a \"Train\" subfolder\nTRAIN_DIR = os.path.join(DATA_DIR, \"Train\")\n\n# Update paths to point to the Train subfolder\nOBO_PATH = os.path.join(TRAIN_DIR, \"go-basic.obo\")\nTRAIN_SEQ_PATH = os.path.join(TRAIN_DIR, \"train_sequences.fasta\")\n\nprint(f\"📂 Reading data from: {TRAIN_DIR}\")\n\n# --- 4. Load Data to Test ---\nprint(\"\\n⏳ Loading Training Sequences (First 3 examples)...\")\nsequences = []\n\n# This should now work correctly\nif os.path.exists(TRAIN_SEQ_PATH):\n    for i, record in enumerate(SeqIO.parse(TRAIN_SEQ_PATH, \"fasta\")):\n        sequences.append(record)\n        if i >= 2: break\n        \n    for seq in sequences:\n        print(f\"   - ID: {seq.id} | Length: {len(seq.seq)}\")\nelse:\n    print(f\"❌ Error: File still not found at {TRAIN_SEQ_PATH}\")\n    # List directory to debug\n    print(f\"Contents of {TRAIN_DIR}:\")\n    print(os.listdir(TRAIN_DIR))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:02:38.708384Z","iopub.execute_input":"2026-01-24T12:02:38.708738Z","iopub.status.idle":"2026-01-24T12:02:41.950978Z","shell.execute_reply.started":"2026-01-24T12:02:38.708705Z","shell.execute_reply":"2026-01-24T12:02:41.950162Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 1. Install Transformers (if needed) ---\n# Kaggle usually has this, but this ensures we have the right version\n!pip install transformers torch --quiet\n\n# --- 2. Import Libraries ---\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\n\n# --- 3. Load the Pre-trained ESM-2 Model ---\n# We use a small version (8 Million parameters) for this demo.\n# For the actual competition, you might switch to 'facebook/esm2_t33_650M_UR50D' later.\nmodel_checkpoint = \"facebook/esm2_t6_8M_UR50D\"\n\nprint(f\"⏳ Downloading model: {model_checkpoint} ...\")\ntokenizer = AutoTokenizer.from_pretrained(model_checkpoint)\nmodel = AutoModel.from_pretrained(model_checkpoint)\n\n# Move to GPU if available\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel = model.to(device)\nprint(f\"✅ Model loaded on: {device}\")\n\n# --- 4. Embed a Single Sequence (Demo) ---\n# Let's take the first sequence you loaded: MOTSC_HUMAN (Length 16)\nseq_id = sequences[0].id\nseq_str = str(sequences[0].seq)\n\nprint(f\"\\n🧪 Processing: {seq_id}\")\nprint(f\"   Sequence: {seq_str}\")\n\n# Tokenize (Convert letters to input IDs)\ninputs = tokenizer(seq_str, return_tensors=\"pt\", add_special_tokens=True).to(device)\n\n# Pass through model\nwith torch.no_grad():\n    outputs = model(**inputs)\n    # The 'last_hidden_state' contains the embedding for every amino acid\n    # Shape: (Batch_Size, Sequence_Length, Embedding_Dimension)\n    embeddings = outputs.last_hidden_state\n\n# We usually take the MEAN of all tokens to get a single vector for the whole protein\nprotein_vector = embeddings.mean(dim=1)\n\nprint(f\"\\n🎉 Result:\")\nprint(f\"   Output Shape: {protein_vector.shape}\")\nprint(f\"   First 5 values: {protein_vector[0, :5].tolist()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:03:21.763055Z","iopub.execute_input":"2026-01-24T12:03:21.76383Z","iopub.status.idle":"2026-01-24T12:03:55.038266Z","shell.execute_reply.started":"2026-01-24T12:03:21.763793Z","shell.execute_reply":"2026-01-24T12:03:55.03749Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom tqdm.notebook import tqdm\nimport numpy as np\nimport pickle\n\n# --- Configuration ---\nBATCH_SIZE = 64  # Adjust if you run out of GPU memory\nMAX_LEN = 1024   # Truncate extremely long proteins to avoid errors\n\ndef get_embeddings(sequences, model, tokenizer, batch_size=32):\n    model.eval()\n    embeddings = {}\n    \n    # Create batches\n    for i in tqdm(range(0, len(sequences), batch_size), desc=\"Processing Batches\"):\n        batch_records = sequences[i:i + batch_size]\n        batch_ids = [rec.id.split('|')[1] for rec in batch_records] # Extract ID (e.g., P9WHI7)\n        batch_seqs = [str(rec.seq)[:MAX_LEN] for rec in batch_records] # Truncate if needed\n\n        # Tokenize\n        inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=MAX_LEN)\n        inputs = {k: v.to(device) for k, v in inputs.items()}\n\n        # Model Inference\n        with torch.no_grad():\n            outputs = model(**inputs)\n            \n        # Extract Embeddings (Mean Pooling)\n        # shape: (batch_size, seq_len, hidden_dim) -> (batch_size, hidden_dim)\n        batch_embeddings = outputs.last_hidden_state.mean(dim=1)\n        \n        # Store in dictionary (move to CPU to save GPU RAM)\n        for seq_id, emb in zip(batch_ids, batch_embeddings):\n            embeddings[seq_id] = emb.cpu().numpy()\n            \n    return embeddings\n\n# --- 1. Load All Sequences ---\nprint(\"⏳ Reading all FASTA sequences...\")\nall_sequences = []\n# Re-open the file to ensure we start from the beginning\nfor record in SeqIO.parse(TRAIN_SEQ_PATH, \"fasta\"):\n    all_sequences.append(record)\n\nprint(f\"✅ Found {len(all_sequences)} sequences.\")\n\n# --- 2. Run Batch Processing ---\nprint(f\"🚀 Starting embedding generation on {device}...\")\ntrain_embeddings = get_embeddings(all_sequences, model, tokenizer, BATCH_SIZE)\n\n# --- 3. Save Results ---\n# We save as a pickle file so we can load it easily later\nSAVE_PATH = \"train_embeddings_esm2_t6.pkl\"\nwith open(SAVE_PATH, \"wb\") as f:\n    pickle.dump(train_embeddings, f)\n\nprint(f\"\\n🎉 Done! Embeddings saved to '{SAVE_PATH}'\")\nprint(f\"   Total embeddings: {len(train_embeddings)}\")\nprint(f\"   Embedding dimension: {train_embeddings[list(train_embeddings.keys())[0]].shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:04:59.81094Z","iopub.execute_input":"2026-01-24T12:04:59.812242Z","iopub.status.idle":"2026-01-24T12:30:55.449174Z","shell.execute_reply.started":"2026-01-24T12:04:59.812209Z","shell.execute_reply":"2026-01-24T12:30:55.448318Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\n# --- 1. Load Data ---\nprint(f\"⏳ Reloading {TRAIN_TERMS_PATH}...\")\ndf_terms = pd.read_csv(TRAIN_TERMS_PATH, sep=\"\\t\")\n\n# --- 2. DEBUG: Inspect the data ---\nprint(\"\\n🔍 Column Names:\", df_terms.columns.tolist())\nprint(\"🔍 First 3 rows:\")\nprint(df_terms.head(3))\n\nprint(\"\\n🔍 Unique values in 'aspect' column:\")\nunique_aspects = df_terms['aspect'].unique()\nprint(unique_aspects)\n\n# --- 3. Smart Selection (Auto-Fix) ---\n# We check which value corresponds to Molecular Function\npossible_mfo_codes = ['MFO', 'F', 'molecular_function']\ntarget_code = None\n\nfor code in possible_mfo_codes:\n    if code in unique_aspects:\n        target_code = code\n        break\n\nif target_code:\n    print(f\"\\n✅ Found correct Aspect Code: '{target_code}'\")\n    \n    # --- 4. Resume Filtering ---\n    df_mfo = df_terms[df_terms['aspect'] == target_code].copy()\n    print(f\"   Total MFO annotations: {len(df_mfo)}\")\n\n    # Recalculate top terms\n    NUM_LABELS = 500\n    term_counts = df_mfo['term'].value_counts()\n    selected_terms = term_counts.head(NUM_LABELS).index.tolist()\n\n    print(f\"✅ Selected top {NUM_LABELS} terms.\")\n    print(f\"   Most common: {selected_terms[0]} ({term_counts.iloc[0]} counts)\")\n    \n    # --- 5. Build Matrix ---\n    print(\"\\n⏳ Re-building Label Matrix...\")\n    protein_ids_with_embeddings = list(train_embeddings.keys())\n    \n    df_filtered = df_mfo[\n        (df_mfo['EntryID'].isin(protein_ids_with_embeddings)) & \n        (df_mfo['term'].isin(selected_terms))\n    ]\n    \n    Y_df = df_filtered.pivot_table(\n        index='EntryID', \n        columns='term', \n        aggfunc='size', \n        fill_value=0\n    )\n    Y_df = (Y_df > 0).astype(int)\n    Y_df = Y_df.reindex(protein_ids_with_embeddings, fill_value=0)\n    \n    print(\"🎉 Label Matrix Ready!\")\n    print(f\"   Y Shape: {Y_df.shape}\")\n\nelse:\n    print(\"\\n❌ CRITICAL ERROR: Could not find 'MFO' or 'F' in aspect column.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:33:48.019449Z","iopub.execute_input":"2026-01-24T12:33:48.020416Z","iopub.status.idle":"2026-01-24T12:33:49.163221Z","shell.execute_reply.started":"2026-01-24T12:33:48.020357Z","shell.execute_reply":"2026-01-24T12:33:49.16253Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, TensorDataset\nfrom sklearn.model_selection import train_test_split\nfrom tqdm.notebook import tqdm\n\n# --- 1. Prepare Data for PyTorch ---\nprint(\"⏳ Converting data to Tensors...\")\n\n# Extract X in the exact order of Y's index to ensure alignment\n# (Y_df.index contains the Protein IDs)\nX_list = [train_embeddings[uid] for uid in Y_df.index]\nX_np = np.array(X_list)\nY_np = Y_df.values\n\n# Convert to Tensors\nX_tensor = torch.tensor(X_np, dtype=torch.float32)\nY_tensor = torch.tensor(Y_np, dtype=torch.float32)\n\n# Split into Train (80%) and Validation (20%)\nX_train, X_val, Y_train, Y_val = train_test_split(X_tensor, Y_tensor, test_size=0.2, random_state=42)\n\nprint(f\"✅ Data Split:\")\nprint(f\"   Train: {X_train.shape[0]} samples\")\nprint(f\"   Val:   {X_val.shape[0]} samples\")\n\n# Create DataLoaders (batches the data for training)\nbatch_size = 64\ntrain_dataset = TensorDataset(X_train, Y_train)\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n\n# --- 2. Define the Model Architecture ---\nclass ProteinClassifier(nn.Module):\n    def __init__(self, input_dim, num_classes):\n        super(ProteinClassifier, self).__init__()\n        self.layer1 = nn.Linear(input_dim, 512)\n        self.relu = nn.ReLU()\n        self.dropout = nn.Dropout(0.3)  # Prevents overfitting\n        self.layer2 = nn.Linear(512, num_classes)\n        \n    def forward(self, x):\n        x = self.layer1(x)\n        x = self.relu(x)\n        x = self.dropout(x)\n        x = self.layer2(x) # No sigmoid here! BCEWithLogitsLoss does it for us.\n        return x\n\n# Initialize Model\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel = ProteinClassifier(input_dim=320, num_classes=500).to(device)\nprint(f\"✅ Model initialized on {device}\")\n\n# --- 3. Training Loop ---\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\nEPOCHS = 10\n\nprint(f\"\\n🚀 Starting Training for {EPOCHS} epochs...\")\n\nfor epoch in range(EPOCHS):\n    model.train()\n    running_loss = 0.0\n    \n    # Training Phase\n    for batch_X, batch_Y in train_loader:\n        batch_X, batch_Y = batch_X.to(device), batch_Y.to(device)\n        \n        optimizer.zero_grad()\n        outputs = model(batch_X)\n        loss = criterion(outputs, batch_Y)\n        loss.backward()\n        optimizer.step()\n        \n        running_loss += loss.item()\n        \n    # Validation Phase (Check accuracy on unseen data)\n    model.eval()\n    with torch.no_grad():\n        val_X, val_Y = X_val.to(device), Y_val.to(device)\n        val_outputs = model(val_X)\n        val_loss = criterion(val_outputs, val_Y)\n    \n    print(f\"Epoch [{epoch+1}/{EPOCHS}] \"\n          f\"Train Loss: {running_loss/len(train_loader):.4f} | \"\n          f\"Val Loss: {val_loss.item():.4f}\")\n\nprint(\"🎉 Training Complete!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:34:34.302546Z","iopub.execute_input":"2026-01-24T12:34:34.302868Z","iopub.status.idle":"2026-01-24T12:34:56.5362Z","shell.execute_reply.started":"2026-01-24T12:34:34.30284Z","shell.execute_reply":"2026-01-24T12:34:56.535488Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- 1. Re-Load the ESM-2 Model (The Feature Extractor) ---\nprint(\"⏳ Re-loading ESM-2 model for embedding generation...\")\nfrom transformers import AutoModel\nesm_model = AutoModel.from_pretrained(\"facebook/esm2_t6_8M_UR50D\")\nesm_model = esm_model.to(device)\nesm_model.eval() \n\n# --- 2. Generate Embeddings using ESM-2 ---\n# (Using the 'get_embeddings' function we fixed in the previous step)\nprint(f\"🚀 Generating embeddings for test set...\")\n\n# CRITICAL FIX: We pass 'esm_model' here, NOT 'model' (which is your classifier)\ntest_embeddings = get_embeddings(test_sequences, esm_model, tokenizer, batch_size=64)\n\nprint(f\"✅ Generated {len(test_embeddings)} test embeddings.\")\n\n# --- 3. Run Inference using the Classifier ---\n# Now we switch back to using 'model' (your trained Classifier)\nprint(\"🔮 Running Classifier Predictions...\")\npredictions = []\nmodel.eval() # This is your ProteinClassifier\n\ntest_ids = list(test_embeddings.keys())\npred_batch_size = 1000\n\nfor i in tqdm(range(0, len(test_ids), pred_batch_size), desc=\"Predicting\"):\n    batch_ids = test_ids[i:i+pred_batch_size]\n    \n    # Get embeddings\n    batch_X_np = np.array([test_embeddings[uid] for uid in batch_ids])\n    batch_X_tensor = torch.tensor(batch_X_np, dtype=torch.float32).to(device)\n    \n    # Predict using the Classifier\n    with torch.no_grad():\n        logits = model(batch_X_tensor)\n        probs = torch.sigmoid(logits).cpu().numpy()\n    \n    # Format Output\n    for j, protein_id in enumerate(batch_ids):\n        protein_probs = probs[j]\n        for term_idx, prob in enumerate(protein_probs):\n            if prob > 0.01:\n                predictions.append(f\"{protein_id}\\t{selected_terms[term_idx]}\\t{prob:.3f}\")\n\n# --- 4. Save Submission ---\nwith open(SUBMISSION_FILE, \"w\") as f:\n    f.write(\"\\n\".join(predictions))\n\nprint(f\"🎉 Submission Ready: {SUBMISSION_FILE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-24T12:37:47.941618Z","iopub.execute_input":"2026-01-24T12:37:47.942232Z","iopub.status.idle":"2026-01-24T13:45:09.891104Z","shell.execute_reply.started":"2026-01-24T12:37:47.942201Z","shell.execute_reply":"2026-01-24T13:45:09.890401Z"}},"outputs":[],"execution_count":null}]}