{"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"},{"sourceId":13436485,"sourceType":"datasetVersion","datasetId":8528539}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# **🧩What You’re Predicting**\n\nThe competition aims to build a model that learns patterns from protein sequences and predicts their correct class or property using advanced models like ESM.\n\n✅ Type of Task: Multi-label classification                                           \n✅ Input: Protein amino acid sequence                                                 \n✅ Output: A list of GO terms for each protein","metadata":{}},{"cell_type":"code","source":"!pip install biopython\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:09:56.199437Z","iopub.execute_input":"2025-10-20T16:09:56.199650Z","iopub.status.idle":"2025-10-20T16:10:01.995391Z","shell.execute_reply.started":"2025-10-20T16:09:56.199624Z","shell.execute_reply":"2025-10-20T16:10:01.994595Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from Bio import SeqIO\nimport pandas as pd","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:01.997208Z","iopub.execute_input":"2025-10-20T16:10:01.997438Z","iopub.status.idle":"2025-10-20T16:10:02.314386Z","shell.execute_reply.started":"2025-10-20T16:10:01.997417Z","shell.execute_reply":"2025-10-20T16:10:02.313814Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fasta_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta\"\nsequences = []\nfor record in SeqIO.parse(fasta_path, \"fasta\"):\n    sequences.append({\n        \"protein_id\": record.id.split('|')[1],  # Extract UniProt ID (e.g., P9WHI7)\n        \"sequence\": str(record.seq)\n    })\n\n\ndf_seq = pd.DataFrame(sequences)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:02.315019Z","iopub.execute_input":"2025-10-20T16:10:02.315489Z","iopub.status.idle":"2025-10-20T16:10:03.263620Z","shell.execute_reply.started":"2025-10-20T16:10:02.315466Z","shell.execute_reply":"2025-10-20T16:10:03.263064Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_seq.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:03.264255Z","iopub.execute_input":"2025-10-20T16:10:03.264453Z","iopub.status.idle":"2025-10-20T16:10:03.285345Z","shell.execute_reply.started":"2025-10-20T16:10:03.264437Z","shell.execute_reply":"2025-10-20T16:10:03.284804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_seq.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:03.286123Z","iopub.execute_input":"2025-10-20T16:10:03.286456Z","iopub.status.idle":"2025-10-20T16:10:03.291030Z","shell.execute_reply.started":"2025-10-20T16:10:03.286439Z","shell.execute_reply":"2025-10-20T16:10:03.290322Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"labels_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\"\ndf_lab = pd.read_csv(labels_path, sep=\"\\t\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:03.292918Z","iopub.execute_input":"2025-10-20T16:10:03.293117Z","iopub.status.idle":"2025-10-20T16:10:03.601662Z","shell.execute_reply.started":"2025-10-20T16:10:03.293103Z","shell.execute_reply":"2025-10-20T16:10:03.601044Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_lab.rename(columns={\n    \"EntryID\": \"protein_id\",\n    \"term\": \"go_term\",\n    \"aspect\": \"ontology\"\n}, inplace=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:03.602360Z","iopub.execute_input":"2025-10-20T16:10:03.602547Z","iopub.status.idle":"2025-10-20T16:10:03.608233Z","shell.execute_reply.started":"2025-10-20T16:10:03.602531Z","shell.execute_reply":"2025-10-20T16:10:03.607603Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_lab.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:03.608993Z","iopub.execute_input":"2025-10-20T16:10:03.609214Z","iopub.status.idle":"2025-10-20T16:10:03.621865Z","shell.execute_reply.started":"2025-10-20T16:10:03.609197Z","shell.execute_reply":"2025-10-20T16:10:03.621291Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Group GO terms for each protein\nprotein_go = df_lab.groupby(\"protein_id\")[\"go_term\"].apply(list).reset_index()\n\n# Merge with sequences\ndf = df_seq.merge(protein_go, on=\"protein_id\", how=\"inner\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:03.622782Z","iopub.execute_input":"2025-10-20T16:10:03.623024Z","iopub.status.idle":"2025-10-20T16:10:04.971642Z","shell.execute_reply.started":"2025-10-20T16:10:03.622997Z","shell.execute_reply":"2025-10-20T16:10:04.970852Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:04.972486Z","iopub.execute_input":"2025-10-20T16:10:04.972707Z","iopub.status.idle":"2025-10-20T16:10:04.981128Z","shell.execute_reply.started":"2025-10-20T16:10:04.972687Z","shell.execute_reply":"2025-10-20T16:10:04.980420Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.isnull().sum()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-20T16:10:04.981900Z","iopub.execute_input":"2025-10-20T16:10:04.982258Z","iopub.status.idle":"2025-10-20T16:10:05.014627Z","shell.execute_reply.started":"2025-10-20T16:10:04.982235Z","shell.execute_reply":"2025-10-20T16:10:05.014037Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom transformers import AutoTokenizer, AutoModel\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nimport joblib\nimport os\nimport gc\n\n# Load Model\nmodel_path = \"/kaggle/input/esm2-t6-8m-ur50d\"\ntokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)\nmodel = AutoModel.from_pretrained(model_path, local_files_only=True)\nmodel.eval()\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\n\n# use less data\ndf_sampled = df.sample(n=50000, random_state=42).reset_index(drop=True)\nsequences = df_sampled[\"sequence\"].tolist()\n\n# Helper Function\ndef embed_full_sequence(seq, tokenizer, model, device, max_length=256):\n    \"\"\"Splits long sequences into chunks, embeds each, averages them.\"\"\"\n    embeddings = []\n    for i in range(0, len(seq), max_length):\n        chunk = seq[i:i+max_length]\n        inputs = tokenizer(chunk, return_tensors=\"pt\", truncation=True, padding=True).to(device)\n        with torch.no_grad():\n            outputs = model(**inputs)\n        emb = outputs.last_hidden_state.mean(1).cpu().numpy()\n        embeddings.append(emb)\n        torch.cuda.empty_cache()\n    return np.mean(embeddings, axis=0)\n\n# Save Progressively\nsave_path = \"/kaggle/working/esm2_50k_full_embeds.pkl\"\nembeddings = []\nstart_idx = 0\n\nif os.path.exists(save_path):\n    embeddings = joblib.load(save_path)\n    start_idx = len(embeddings)\n    print(f\"Resuming from index {start_idx} / {len(sequences)}\")\n\nfor idx in tqdm(range(start_idx, len(sequences)), total=len(sequences)):\n    seq = sequences[idx]\n    try:\n        emb = embed_full_sequence(seq, tokenizer, model, device, max_length=256)\n        embeddings.append(emb)\n\n        # Save progress every 100 sequences\n        if idx % 100 == 0:\n            joblib.dump(embeddings, save_path)\n            print(f\"Saved up to index {idx}\")\n\n    except RuntimeError as e:\n        print(f\"Skipped {idx} due to memory error: {e}\")\n        torch.cuda.empty_cache()\n        gc.collect()\n        continue\n\n# Final save\njoblib.dump(embeddings, save_path)\nprint(\"embeddings saved safely.\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}