{"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":14084779,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"List the Data Files","metadata":{}},{"cell_type":"code","source":"import os\n\n# The directory where the competition data is stored\ndata_dir = '/kaggle/input/cafa-6-protein-function-prediction'\n\n# List all the files in that directory\nfiles = os.listdir(data_dir)\nprint(files)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:36.948398Z","iopub.execute_input":"2025-10-19T14:25:36.949246Z","iopub.status.idle":"2025-10-19T14:25:36.968811Z","shell.execute_reply.started":"2025-10-19T14:25:36.949214Z","shell.execute_reply":"2025-10-19T14:25:36.967992Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Load the Training Labels¶","metadata":{}},{"cell_type":"code","source":"import pandas as pd\n\n# Define the base directory for the competition data\ndata_dir = '/kaggle/input/cafa-6-protein-function-prediction'\n\n# The training files are inside the 'Train' subdirectory\ntrain_dir = os.path.join(data_dir, 'Train')\n\n# Define the full, correct path to the training terms file\ntrain_terms_path = os.path.join(train_dir, 'train_terms.tsv')\n\nprint(f\"Attempting to load file from: {train_terms_path}\")\n\n# Load the data\ntrain_df = pd.read_csv(train_terms_path, sep='\\t')\n\n# Display the size and first 5 rows to confirm it loaded\nprint(f\"\\nSuccess! Shape of the training labels DataFrame: {train_df.shape}\")\ndisplay(train_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:36.969959Z","iopub.execute_input":"2025-10-19T14:25:36.970210Z","iopub.status.idle":"2025-10-19T14:25:38.095225Z","shell.execute_reply.started":"2025-10-19T14:25:36.970191Z","shell.execute_reply":"2025-10-19T14:25:38.094539Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Load the Protein Sequence","metadata":{}},{"cell_type":"code","source":"def load_sequences_corrected(filepath):\n    \"\"\"\n    Reads a FASTA file, correctly parsing the headers to extract the main protein ID,\n    and returns a dictionary mapping these IDs to sequences.\n    \"\"\"\n    sequences = {}\n    current_id = \"\"\n    with open(filepath, 'r') as f:\n        for line in f:\n            line = line.strip()\n            if line.startswith('>'):\n                # Split the header by '|' and take the second element\n                # e.g., from '>sp|Q9Y2X8|CD3E_HUMAN' we get 'Q9Y2X8'\n                current_id = line.split('|')[1] \n                sequences[current_id] = \"\"\n            else:\n                sequences[current_id] += line\n    return sequences\n\n# --- RERUN THE LOADING AND ANALYSIS ---\n\n# Define the file path for the sequences\ntrain_seq_path = os.path.join(train_dir, 'train_sequences.fasta')\n\n# Load the sequences using our NEW function\ntrain_sequences = load_sequences_corrected(train_seq_path)\n\n# Now, rerun the intersection check\nlabel_protein_ids = set(train_df['EntryID'].unique())\nsequence_protein_ids = set(train_sequences.keys())\ncommon_protein_ids = label_protein_ids.intersection(sequence_protein_ids)\n\nprint(f\"Number of proteins with labels: {len(label_protein_ids)}\")\nprint(f\"Number of proteins with sequences: {len(sequence_protein_ids)}\")\nprint(f\"Number of proteins we can use for training (in both sets): {len(common_protein_ids)}\")\n\n# --- Safely print an example ---\nif common_protein_ids:\n    example_protein_id = list(common_protein_ids)[0]\n    \n    print(f\"\\nExample of a common protein ID: '{example_protein_id}'\")\n    print(f\"Sequence for this protein:\")\n    print(train_sequences[example_protein_id])\nelse:\n    print(\"\\nWarning: Still no common proteins found. The header format might be different.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:38.096090Z","iopub.execute_input":"2025-10-19T14:25:38.096381Z","iopub.status.idle":"2025-10-19T14:25:38.904240Z","shell.execute_reply.started":"2025-10-19T14:25:38.096348Z","shell.execute_reply":"2025-10-19T14:25:38.903513Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\n\n# --- Step 1: Filter the DataFrame ---\n# Keep only the labels for the proteins that we have sequences for.\ntrain_df_filtered = train_df[train_df['EntryID'].isin(common_protein_ids)].copy()\n\nprint(f\"Original number of labels: {len(train_df)}\")\nprint(f\"Number of labels for common proteins: {len(train_df_filtered)}\")\n\n# --- Step 2: Identify the Top 1,500 GO Terms ---\n# Count the occurrences of each term\nterm_counts = train_df_filtered['term'].value_counts()\n# Get the names of the top 1500 terms\ntop_terms = term_counts.head(1500).index.tolist()\n\nprint(f\"\\nTotal unique GO terms: {len(term_counts)}\")\nprint(f\"We will predict the top {len(top_terms)} most frequent terms.\")\n\n# --- Step 3: Filter the DataFrame Again ---\n# Keep only the rows that correspond to one of the top 1500 terms.\ntrain_df_top = train_df_filtered[train_df_filtered['term'].isin(top_terms)].copy()\n\n\n# --- Step 4: Create the Binary Label Matrix (the \"wide\" format) ---\n# Add a 'value' column of 1s to help with the pivot operation\ntrain_df_top['value'] = 1\n\n# Create the pivot table\n# Rows: proteins (EntryID), Columns: functions (term), Values: 1 if present, 0 otherwise\nlabels_df = train_df_top.pivot_table(\n    index='EntryID', \n    columns='term', \n    values='value', \n    fill_value=0\n)\n\n# --- Final Check ---\nprint(f\"\\nShape of our final labels DataFrame: {labels_df.shape}\")\nprint(\"This means we have labels for\", labels_df.shape[0], \"proteins and\", labels_df.shape[1], \"unique functions.\")\ndisplay(labels_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:38.905504Z","iopub.execute_input":"2025-10-19T14:25:38.905768Z","iopub.status.idle":"2025-10-19T14:25:41.939924Z","shell.execute_reply.started":"2025-10-19T14:25:38.905749Z","shell.execute_reply":"2025-10-19T14:25:41.939014Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Prepare the Labels for the Model","metadata":{}},{"cell_type":"code","source":"import numpy as np\n\n# --- Step 4.1: Filter the DataFrame ---\n# We only want to work with labels for the proteins that we actually have sequences for.\ntrain_df_filtered = train_df[train_df['EntryID'].isin(common_protein_ids)].copy()\n\nprint(f\"Original number of labels: {len(train_df)}\")\nprint(f\"Number of labels for common proteins: {len(train_df_filtered)}\")\n\n\n# --- Step 4.2: Identify the Top 1,500 GO Terms ---\n# First, we count how many times each GO term appears.\nterm_counts = train_df_filtered['term'].value_counts()\n# Then, we get the names (the IDs) of the 1500 most frequent terms.\ntop_terms = term_counts.head(1500).index.tolist()\n\nprint(f\"\\nTotal unique GO terms: {len(term_counts)}\")\nprint(f\"We will focus on predicting the top {len(top_terms)} most frequent terms.\")\n\n\n# --- Step 4.3: Filter Again for Top Terms ---\n# Now we create a new DataFrame that only includes the labels for our top terms.\ntrain_df_top = train_df_filtered[train_df_filtered['term'].isin(top_terms)].copy()\n\n\n# --- Step 4.4: Create the Binary Label Matrix ---\n# **FIX:** Add a 'value' column containing all 1s. This is what pivot_table will use.\ntrain_df_top['value'] = 1\n\n# Now, create the pivot table.\n# It uses proteins as the index (rows) and terms as the columns.\nlabels_df = train_df_top.pivot_table(\n    index='EntryID',\n    columns='term',\n    values='value',\n    fill_value=0\n)\n\n\n# --- Step 4.5: Final Check ---\nprint(f\"\\nShape of our final labels DataFrame: {labels_df.shape}\")\nprint(f\"This matrix has {labels_df.shape[0]} proteins and {labels_df.shape[1]} functions.\")\nprint(\"\\nHere's a preview of the final label matrix:\")\ndisplay(labels_df.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:41.940850Z","iopub.execute_input":"2025-10-19T14:25:41.941113Z","iopub.status.idle":"2025-10-19T14:25:44.884877Z","shell.execute_reply.started":"2025-10-19T14:25:41.941088Z","shell.execute_reply":"2025-10-19T14:25:44.884192Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Rerunning the logic from CELL 3 and CELL 3.5 to redefine lost variables.\n# This fixes the NameError.\n\n# 1. Rerun CELL 3's core logic (Define initial labels_df)\ntrain_df_filtered = train_df[train_df['EntryID'].isin(common_protein_ids)].copy()\nterm_counts = train_df_filtered['term'].value_counts()\ntop_terms = term_counts.head(1500).index.tolist()\ntrain_df_top = train_df_filtered[train_df_filtered['term'].isin(top_terms)].copy()\ntrain_df_top['value'] = 1\nlabels_df = train_df_top.pivot_table(index='EntryID', columns='term', values='value', fill_value=0)\n\n# 2. Rerun CELL 3.5's core logic (Filter and define final_protein_ids)\nfinal_protein_ids = labels_df.index.tolist()\nMAX_SEQUENCE_LENGTH = 1000 # Using the established safe limit\n\nids_to_process_filtered = []\nfor pid in final_protein_ids:\n    sequence = train_sequences.get(pid)\n    if sequence and len(sequence) <= MAX_SEQUENCE_LENGTH:\n        ids_to_process_filtered.append(pid)\n\nfinal_protein_ids = ids_to_process_filtered\nlabels_df = labels_df.loc[final_protein_ids] # Re-filter the labels DataFrame\n\nprint(\"final_protein_ids has been redefined successfully. You can now run CELL 4.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:44.885681Z","iopub.execute_input":"2025-10-19T14:25:44.885934Z","iopub.status.idle":"2025-10-19T14:25:48.190218Z","shell.execute_reply.started":"2025-10-19T14:25:44.885917Z","shell.execute_reply":"2025-10-19T14:25:48.189279Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Step 5.1: Install and Import Libraries ---\n# We need the 'transformers' library from Hugging Face to load the ESM-2 model.\n!pip install transformers\n\nimport torch\nfrom transformers import EsmModel, EsmTokenizer\nfrom tqdm import tqdm\nimport numpy as np\n\n# --- Step 5.2: Set Up the Model ---\n# Check if a GPU is available and set the device accordingly.\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# Load the pre-trained ESM-2 model and its tokenizer.\n# We're using a smaller version to be efficient with our GPU quota.\nmodel_name = \"facebook/esm2_t30_150M_UR50D\"\ntokenizer = EsmTokenizer.from_pretrained(model_name)\nmodel = EsmModel.from_pretrained(model_name).to(device)\nmodel.eval() # Set the model to evaluation mode\n\n# --- Step 5.3: Generate Embeddings ---\n# We will store the embeddings in this dictionary.\nprotein_embeddings = {}\n\n# Get the list of protein IDs we need to process (the ones in our label matrix).\nids_to_process = labels_df.index.tolist()\n\n# Let's start with a small sample to test our code.\n# Change this to len(ids_to_process) to run on the full dataset later.\nsample_size = 100 \n\nprint(f\"\\nGenerating embeddings for the first {sample_size} proteins...\")\n\n# Use tqdm for a nice progress bar.\nfor protein_id in tqdm(ids_to_process[:sample_size]):\n    sequence = train_sequences[protein_id]\n    \n    # Tokenize the sequence\n    inputs = tokenizer(sequence, return_tensors=\"pt\", add_special_tokens=False).to(device)\n    \n    # Get the model's output (no gradient calculation needed)\n    with torch.no_grad():\n        outputs = model(**inputs)\n\n    # The model outputs embeddings for each amino acid. We'll take the mean\n    # across the sequence length to get a single embedding for the whole protein.\n    embedding = outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()\n    protein_embeddings[protein_id] = embedding\n\n# --- Step 5.4: Final Check ---\nprint(f\"\\nSuccessfully generated {len(protein_embeddings)} embeddings.\")\nexample_id = list(protein_embeddings.keys())[0]\nprint(f\"The embedding for protein '{example_id}' has a shape of: {protein_embeddings[example_id].shape}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:25:48.191244Z","iopub.execute_input":"2025-10-19T14:25:48.191793Z","iopub.status.idle":"2025-10-19T14:26:22.894417Z","shell.execute_reply.started":"2025-10-19T14:25:48.191765Z","shell.execute_reply":"2025-10-19T14:26:22.893529Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nfrom transformers import EsmModel, EsmTokenizer\nfrom tqdm import tqdm\nimport numpy as np\nimport gc # Garbage Collector\n\n# --- Step 6.1: Set Up the Model ---\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\nmodel_name = \"facebook/esm2_t30_150M_UR50D\"\ntokenizer = EsmTokenizer.from_pretrained(model_name)\nmodel = EsmModel.from_pretrained(model_name).to(device)\nmodel.eval()\n\n# --- Step 6.2: Make the Process Resumable ---\nembeddings_file = '/kaggle/working/protein_embeddings.npy'\nprotein_embeddings = {}\n\nif os.path.exists(embeddings_file):\n    print(\"Found existing embeddings file. Loading to resume progress...\")\n    # allow_pickle=True is necessary for loading a dictionary\n    protein_embeddings = np.load(embeddings_file, allow_pickle=True).item()\n    print(f\"Resuming with {len(protein_embeddings)} embeddings already generated.\")\n\n# Get the list of protein IDs we need to process\nall_ids = labels_df.index.tolist()\n# Filter out the IDs we've already processed\nids_to_process = [pid for pid in all_ids if pid not in protein_embeddings]\nprint(f\"Number of new proteins to process: {len(ids_to_process)}\")\n\n\n# --- Step 6.3: Generate Embeddings in Batches ---\nbatch_size = 16 # Process 16 proteins at a time\nprint(f\"Generating embeddings with a batch size of {batch_size}...\")\n\nfor i in tqdm(range(0, len(ids_to_process), batch_size)):\n    batch_ids = ids_to_process[i:i+batch_size]\n    sequences = [train_sequences[pid] for pid in batch_ids]\n    \n    # **THE FIX:** Tokenize with truncation to prevent out-of-memory errors\n    inputs = tokenizer(\n        sequences, \n        return_tensors=\"pt\", \n        padding=True,          # Pad sequences to the same length within a batch\n        truncation=True,       # Truncate sequences longer than the model's max length\n        max_length=1022,       # A safe max length for ESM-2\n        add_special_tokens=False\n    ).to(device)\n    \n    with torch.no_grad():\n        outputs = model(**inputs)\n\n    # Extract embeddings and store them\n    for j, protein_id in enumerate(batch_ids):\n        # We take the mean of the embeddings for all amino acids in the sequence\n        embedding = outputs.last_hidden_state[j].mean(dim=0).cpu().numpy()\n        protein_embeddings[protein_id] = embedding\n    \n    # Periodically save our progress\n    if i % (batch_size * 100) == 0: # Save every 100 batches\n        np.save(embeddings_file, protein_embeddings)\n        \n    # Clean up GPU memory\n    del inputs, outputs\n    gc.collect()\n    torch.cuda.empty_cache()\n\n\n# --- Step 6.4: Final Save ---\nprint(f\"\\nSuccessfully generated embeddings for {len(protein_embeddings)} total proteins.\")\nnp.save(embeddings_file, protein_embeddings)\nprint(f\"\\nAll embeddings have been saved to '{embeddings_file}'.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-19T14:27:30.162352Z","iopub.execute_input":"2025-10-19T14:27:30.163109Z"}},"outputs":[],"execution_count":null}]}