{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =================================================================\n# NOTEBOOK 2: TEST EMBEDDING GENERATION\n# Purpose: Generate and save ONLY the test embeddings.\n# =================================================================\n!pip install biopython transformers[torch]\n\nimport torch\nfrom transformers import AutoTokenizer, EsmModel\nfrom Bio import SeqIO\nimport numpy as np\nfrom tqdm import tqdm\nimport gc\n\n# --- Setup Model and Device ---\nprint(\"Setting up the model and tokenizer...\")\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\nmodel_name = \"facebook/esm2_t6_8M_UR50D\" \ntokenizer = AutoTokenizer.from_pretrained(model_name)\nmodel = EsmModel.from_pretrained(model_name).to(device)\nmodel.eval()\n\ndef get_embedding(sequence: str) -> np.ndarray:\n    inputs = tokenizer(sequence, return_tensors='pt', truncation=True, max_length=1022).to(device)\n    with torch.no_grad():\n        outputs = model(**inputs)\n    embedding = outputs.last_hidden_state[0, 1:-1].mean(dim=0).cpu().numpy()\n    return embedding\n\n# --- Process and Save TEST Embeddings ---\nprint(\"\\nProcessing test (superset) sequences...\")\ntest_sequences_file = '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\ntest_embeddings = {}\n\nfor record in tqdm(SeqIO.parse(test_sequences_file, \"fasta\"), desc=\"Generating Test Embeddings\"):\n    protein_id = record.id\n    sequence = str(record.seq)\n    test_embeddings[protein_id] = get_embedding(sequence)\n\nprint(f\"Generated {len(test_embeddings)} embeddings for the test set.\")\nprint(\"Saving test embeddings to /kaggle/working/test_embeddings.npz...\")\nnp.savez_compressed('/kaggle/working/test_embeddings.npz', **test_embeddings)\n\nprint(\"\\n🎉 Test embeddings notebook finished successfully!\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}