{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# --------------------------------------------------------------------------\n# NOTEBOOK 1: EMBEDDING GENERATION\n#\n# PURPOSE:\n# This notebook's ONLY job is to generate and save protein embeddings.\n# It reads the FASTA files, runs each sequence through the ESM2 model,\n# and saves the resulting numerical vectors to output files.\n# This notebook will take a long time to run. Commit it and let it finish.\n# --------------------------------------------------------------------------\n\n# --------------------------------------------------------------------------\n# STEP 1: INSTALL REQUIRED LIBRARIES\n# We must install biopython as it is not included in the Kaggle environment by default.\n# --------------------------------------------------------------------------\n!pip install biopython\n!pip install transformers[torch]\n\nprint(\"✅ Libraries installed. Proceeding with the script.\")\n\nimport torch\nfrom transformers import AutoTokenizer, EsmModel\nfrom Bio import SeqIO\nimport numpy as np\nfrom tqdm import tqdm\nimport gc\n\n# ----------------------------------\n# 1. Setup Model and Device\n# ----------------------------------\nprint(\"Setting up the model and tokenizer...\")\n# Setup device (GPU if available)\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# Load pre-trained model and tokenizer\nmodel_name = \"facebook/esm2_t6_8M_UR50D\" \ntokenizer = AutoTokenizer.from_pretrained(model_name)\nmodel = EsmModel.from_pretrained(model_name).to(device)\nmodel.eval() # Set to evaluation mode\n\ndef get_embedding(sequence: str) -> np.ndarray:\n    \"\"\"Generates a protein embedding for a given amino acid sequence.\"\"\"\n    # Truncation is important for long sequences\n    inputs = tokenizer(sequence, return_tensors='pt', truncation=True, max_length=1022).to(device)\n    with torch.no_grad():\n        outputs = model(**inputs)\n    # Take the mean of the last hidden state, ignoring special tokens [CLS] and [SEP]\n    embedding = outputs.last_hidden_state[0, 1:-1].mean(dim=0).cpu().numpy()\n    return embedding\n\n# ----------------------------------\n# 2. Process and Save TRAINING Embeddings\n# ----------------------------------\nprint(\"\\nProcessing training sequences...\")\ntrain_sequences_file = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\ntrain_embeddings = {}\n\n# Use SeqIO.parse for memory-efficient iteration\nfor record in tqdm(SeqIO.parse(train_sequences_file, \"fasta\"), desc=\"Generating Train Embeddings\"):\n    # The FASTA header format for train is \"sp|PROTEIN_ID|GENE_NAME\"\n    protein_id = record.id.split('|')[1] \n    sequence = str(record.seq)\n    train_embeddings[protein_id] = get_embedding(sequence)\n\nprint(f\"Generated {len(train_embeddings)} embeddings for the training set.\")\n\n# Save the training embeddings to a compressed .npz file in the working directory\nprint(\"Saving training embeddings to /kaggle/working/train_embeddings.npz...\")\nnp.savez_compressed('/kaggle/working/train_embeddings.npz', **train_embeddings)\nprint(\"✅ Training embeddings saved.\")\nprint(\"\\n🎉 Training embeddings notebook finished successfully!\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}