{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "3f87122b", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "================================================================================\n", "... 0. PEFT LIBRARY INITIALIZED ...\n", " Ready for Parameter-Efficient Fine-Tuning with LoRA!\n", "šŸ§¬ BALM Adaptation for ESM-2 with LoRA Fine-Tuning\n", "āš” Fast Training with 5-Fold Cross-Validation\n", "================================================================================\n", "šŸ”’ Setting up reproducibility with seed 42\n", "āœ… Reproducibility setup complete\n", "\n", "================================================================================\n", "... 3. LOADING DATA FROM Data.csv ...\n", "Successfully loaded Data.csv with 12062 rows.\n", "\n", "--- Initial Data Check ---\n", " Target \\\n", "0 PKFTKCRSPERETFSCHWTLGPIQLFYTRRNTQEWTQEWKECPDYV... \n", "1 QDNSRYTHFLTQHYDAKPQGRDDRYCESIMRRRGLTSPCKDINTFI... \n", "2 KSFPEVVGKTVDQAREYFTLHYPQYDVYFLPEGSPVTLDLRYNRVR... \n", "3 TNTVAAYNLTWKSTNFKTILEWEPKPVNQVYTVQISTKSGDWKSKC... \n", "4 PIVQNLQGQMVHQAISPRTLNAWVKVVEEKAFSPEVIPMFSALSEG... \n", "\n", " proteina Y PDB \\\n", "0 FPTIPLSRLFDNAMLRAHRLHQLAFDTYQEFEEAYIPKEQKYSFLQ... 9.045757 1A22 \n", "1 SLDIQSLDIQCEELSDARWAELLPLLQQCQVVRLDDCGLTEARCKD... 15.301030 1A4Y \n", "2 CGVPAIQPVLSGLIVNGEEAVPGSWPWQVSLQDKTGFHFCGGSLIN... 11.826814 1ACB \n", "3 DIKMTQSPSSMYASLGERVTITCKASQDIRKYLNWYQQKPWKSPKT... 8.468521 1AHW \n", "4 MVNPTVFFDIAVDGEPLGRVSFELFADKVPKTAENFRALSTGEKGF... 4.920819 1AK4 \n", "\n", " Subgroup SourceDataSet \n", "0 others SKEMPI v2.0 \n", "1 others SKEMPI v2.0 \n", "2 others SKEMPI v2.0 \n", "3 others SKEMPI v2.0 \n", "4 others SKEMPI v2.0 \n", "\n", "Missing values per column:\n", " Target 39\n", "proteina 28\n", "Y 0\n", "PDB 0\n", "Subgroup 0\n", "SourceDataSet 0\n", "dtype: int64\n", "\n", "Warning: Found 67 total missing values. Dropping rows with NaNs.\n", "Dataset now has 12019 rows after cleaning.\n", "\n", "āœ… Data is loaded and cleaned. Final dataset size: 12019 samples.\n", "\n", "Index: 12019 entries, 0 to 12061\n", "Data columns (total 6 columns):\n", " # Column Non-Null Count Dtype \n", "--- ------ -------------- ----- \n", " 0 Target 12019 non-null object \n", " 1 proteina 12019 non-null object \n", " 2 Y 12019 non-null float64\n", " 3 PDB 12019 non-null object \n", " 4 Subgroup 12019 non-null object \n", " 5 SourceDataSet 12019 non-null object \n", "dtypes: float64(1), object(5)\n", "memory usage: 657.3+ KB\n", "\n", "================================================================================\n", "... 4. CONFIGURING THE EXPERIMENT ...\n", "Results will be saved in: cv_results_lora_seqsim\n", "āœ… Configuration set. Using CV Strategy: SequenceSimilarity\n", " Epochs: 30, Patience: 15\n", " LoRA Enabled: Rank=8, Alpha=16, Dropout=0.1\n", "\n", "================================================================================\n", "... 5. DEFINING ADAPTED MODEL AND DATASET ...\n", "\n", "================================================================================\n", "... 6. DEFINING METRICS, EVALUATION & SEQUENCE HELPERS ...\n", "\n", "================================================================================\n", "... 7. RUNNING 'SequenceSimilarity' STRATEGY WITH LoRA (80/20 Split with Training Loss Patience) ...\n", "Global pKd bounds for scaling: [1.32, 15.70]\n", "\n", "šŸ§¬ Step 1: Extracting unique protein sequences...\n", "Found 11598 unique protein sequences.\n", "\n", "šŸ§¬ Step 2: Computing sequence similarity matrix...\n", "Computing 3-mer Jaccard similarity matrix for 11598 sequences...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Computing similarities: 0%| | 13/11598 [00:08<1:59:00, 1.62it/s]\n" ] }, { "ename": "KeyboardInterrupt", "evalue": "", "output_type": "error", "traceback": [ "\u001b[31m---------------------------------------------------------------------------\u001b[39m", "\u001b[31mKeyboardInterrupt\u001b[39m Traceback (most recent call last)", "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 757\u001b[39m\n\u001b[32m 754\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mFold model checkpoints and prediction files are saved in the \u001b[39m\u001b[33m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mRESULTS_DIR\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m'\u001b[39m\u001b[33m directory.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 756\u001b[39m \u001b[38;5;66;03m# --- Call the new CV function ---\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m757\u001b[39m \u001b[43mrun_sequence_similarity_cv_with_lora\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 760\u001b[39m \u001b[38;5;66;03m# ==============================================================================\u001b[39;00m\n\u001b[32m 761\u001b[39m \u001b[38;5;66;03m# 8. FINAL NOTES AND RECOMMENDATIONS\u001b[39;00m\n\u001b[32m 762\u001b[39m \u001b[38;5;66;03m# ==============================================================================\u001b[39;00m\n\u001b[32m 763\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[33m\"\u001b[39m + \u001b[33m\"\u001b[39m\u001b[33m=\u001b[39m\u001b[33m\"\u001b[39m*\u001b[32m80\u001b[39m)\n", "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 551\u001b[39m, in \u001b[36mrun_sequence_similarity_cv_with_lora\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m 548\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mFound \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(unique_sequences)\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m unique protein sequences.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 550\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[33mšŸ§¬ Step 2: Computing sequence similarity matrix...\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m--> \u001b[39m\u001b[32m551\u001b[39m similarity_matrix = \u001b[43mcompute_kmer_similarity_matrix\u001b[49m\u001b[43m(\u001b[49m\u001b[43munique_sequences\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mk\u001b[49m\u001b[43m=\u001b[49m\u001b[43mK_MER_SIZE\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 553\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33m\"\u001b[39m\u001b[38;5;130;01m\\n\u001b[39;00m\u001b[33mšŸ§¬ Step 3: Clustering sequences...\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m 554\u001b[39m clusters = cluster_sequences_by_similarity(similarity_matrix, threshold=SIMILARITY_THRESHOLD, method=CLUSTERING_METHOD)\n", "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 457\u001b[39m, in \u001b[36mcompute_kmer_similarity_matrix\u001b[39m\u001b[34m(sequences, k)\u001b[39m\n\u001b[32m 455\u001b[39m similarity_matrix[i, j] = \u001b[32m1.0\u001b[39m\n\u001b[32m 456\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m457\u001b[39m sim = \u001b[43mjaccard_similarity\u001b[49m\u001b[43m(\u001b[49m\u001b[43msequences\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msequences\u001b[49m\u001b[43m[\u001b[49m\u001b[43mj\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m 458\u001b[39m similarity_matrix[i, j] = sim\n\u001b[32m 459\u001b[39m similarity_matrix[j, i] = sim\n", "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 438\u001b[39m, in \u001b[36mcompute_kmer_similarity_matrix..jaccard_similarity\u001b[39m\u001b[34m(seq1, seq2)\u001b[39m\n\u001b[32m 435\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m {sequence}\n\u001b[32m 436\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m {sequence[i:i+k_size] \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(sequence) - k_size + \u001b[32m1\u001b[39m)}\n\u001b[32m--> \u001b[39m\u001b[32m438\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mjaccard_similarity\u001b[39m(seq1, seq2):\n\u001b[32m 439\u001b[39m \u001b[38;5;250m \u001b[39m\u001b[33;03m\"\"\"Compute Jaccard similarity between two sequences based on k-mers.\"\"\"\u001b[39;00m\n\u001b[32m 440\u001b[39m kmers1 = get_kmers(seq1, k)\n", "\u001b[31mKeyboardInterrupt\u001b[39m: " ] } ], "source": [ "\n", "import os\n", "os.environ.pop('MPLBACKEND', None)\n", "\n", "os.environ['MPLBACKEND'] = 'agg'\n", "\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "\n", "\n", "# ==============================================================================\n", "# 0. PEFT IMPORTS AND INITIAL SETUP (NEW)\n", "# ==============================================================================\n", "# Ensure you have the peft library installed: pip install peft\n", "from peft import LoraConfig, get_peft_model, TaskType\n", "\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 0. PEFT LIBRARY INITIALIZED ...\")\n", "print(\" Ready for Parameter-Efficient Fine-Tuning with LoRA!\")\n", "# ==============================================================================\n", "# 1. IMPORTS AND INITIAL SETUP\n", "# ==============================================================================\n", "import os\n", "import sys\n", "import time\n", "import json\n", "import random\n", "import pickle\n", "from pathlib import Path\n", "from datetime import datetime\n", "\n", "# Standard data science libraries\n", "import numpy as np\n", "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "# In the IMPORTS AND INITIAL SETUP section\n", "from tqdm import tqdm\n", "import warnings\n", "\n", "# Machine learning libraries\n", "from sklearn.model_selection import KFold, GroupKFold\n", "from sklearn.metrics import mean_squared_error\n", "from scipy.stats import pearsonr, spearmanr\n", "from sklearn.cluster import AgglomerativeClustering\n", "\n", "\n", "# PyTorch and Transformers\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torch.utils.data import Dataset, DataLoader\n", "from torch.optim import AdamW\n", "from transformers import AutoModel, AutoTokenizer\n", "\n", "# Advanced CV imports\n", "from Bio.Align import PairwiseAligner\n", "from scipy.cluster.hierarchy import linkage, fcluster\n", "from multiprocessing import Pool, cpu_count\n", "from typing import List, Tuple\n", "\n", "# Set style for plots\n", "\n", "\n", "# --- MODIFIED: Title reflects the use of ESM-2 ---\n", "print(\"šŸ§¬ BALM Adaptation for ESM-2 with LoRA Fine-Tuning\")\n", "print(\"āš” Fast Training with 5-Fold Cross-Validation\")\n", "print(\"=\"*80)\n", "\n", "def setup_reproducibility(seed=42):\n", " \"\"\"Setup complete reproducibility.\"\"\"\n", " print(f\"šŸ”’ Setting up reproducibility with seed {seed}\")\n", " random.seed(seed)\n", " np.random.seed(seed)\n", " torch.manual_seed(seed)\n", " if torch.cuda.is_available():\n", " torch.cuda.manual_seed(seed)\n", " torch.cuda.manual_seed_all(seed)\n", " torch.backends.cudnn.deterministic = True\n", " torch.backends.cudnn.benchmark = False\n", " os.environ['PYTHONHASHSEED'] = str(seed)\n", " print(f\"āœ… Reproducibility setup complete\")\n", "\n", "setup_reproducibility(42)\n", "\n", "# ==============================================================================\n", "# 2. ESM-2 EMBEDDING EXTRACTOR (Modified for LoRA)\n", "# ==============================================================================\n", "\n", "class ProteinEmbeddingExtractor:\n", " \"\"\"\n", " A class to prepare the ESM-2 model, optionally with LoRA adapters.\n", " This class handles model loading and tokenization.\n", " \"\"\"\n", " def __init__(\n", " self,\n", " model_name=\"facebook/esm2_t33_650M_UR50D\",\n", " device=\"auto\",\n", " lora_rank=8,\n", " lora_alpha=16,\n", " lora_dropout=0.1,\n", " use_lora=True\n", " ):\n", " self.model_name = model_name\n", " self.use_lora = use_lora\n", " \n", " if device == \"auto\":\n", " self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " else:\n", " self.device = torch.device(device)\n", " \n", " print(f\"šŸ”§ Initializing ProteinEmbeddingExtractor on {self.device}\")\n", " \n", " self.dtype = torch.float16 if self.device.type == 'cuda' else torch.float32\n", " print(f\"šŸ“„ Loading ESM-2 model: {model_name} with dtype: {self.dtype}\")\n", " \n", " self.model = AutoModel.from_pretrained(\n", " model_name,\n", " torch_dtype=self.dtype\n", " )\n", " self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n", " \n", " if self.use_lora:\n", " print(f\"āœØ Applying LoRA with rank={lora_rank}, alpha={lora_alpha}, dropout={lora_dropout}\")\n", " lora_config = LoraConfig(\n", " r=lora_rank,\n", " lora_alpha=lora_alpha,\n", " lora_dropout=lora_dropout,\n", " bias=\"none\",\n", " task_type=TaskType.FEATURE_EXTRACTION,\n", " target_modules=[\"key\", \"query\", \"value\"]\n", " )\n", " self.model = get_peft_model(self.model, lora_config)\n", " print(\" āœ… LoRA model loaded and configured. Trainable parameters:\")\n", " self.model.print_trainable_parameters()\n", " else:\n", " print(\" āš ļø LoRA is disabled. All base model weights are frozen.\")\n", " for param in self.model.parameters():\n", " param.requires_grad = False\n", "\n", " self.model.to(self.device)\n", " self.embedding_size = self.model.config.hidden_size\n", " print(f\" āœ… Model loaded. Embedding size from config: {self.embedding_size}\")\n", "\n", " def get_model_and_tokenizer(self):\n", " \"\"\"Returns the prepared model (PeftModel or base) and tokenizer.\"\"\"\n", " return self.model, self.tokenizer\n", "\n", "# ==============================================================================\n", "# 3. LOAD YOUR DATASET\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 3. LOADING DATA FROM Data.csv ...\")\n", "\n", "protein1_col = 'Target'\n", "protein2_col = 'proteina'\n", "target_col = 'Y'\n", "pdb_col_raw = 'PDB'\n", "subgroup_col = 'Subgroup'\n", "source_dataset_col = 'Source Data Set' # NEW\n", "try:\n", " df_raw = pd.read_csv(r\"C:\\Users\\hs494\\PPIBALM\\BALM-PPI\\scripts\\notebooks\\PPB_Affinity_Sequences_Final (version 1).csv\")\n", " print(f\"Successfully loaded Data.csv with {len(df_raw)} rows.\")\n", " df = df_raw[[protein1_col, protein2_col, target_col, pdb_col_raw, subgroup_col,source_dataset_col]].copy()\n", " df.rename(columns={\n", " protein1_col: 'Target',\n", " protein2_col: 'proteina',\n", " target_col: 'Y',\n", " subgroup_col : 'Subgroup',\n", " source_dataset_col: 'SourceDataSet' # NEW\n", " }, inplace=True)\n", " \n", " print(\"\\n--- Initial Data Check ---\")\n", " print(df.head())\n", " missing_values = df.isnull().sum()\n", " print(\"\\nMissing values per column:\\n\", missing_values)\n", " \n", " if missing_values.sum() > 0:\n", " print(f\"\\nWarning: Found {missing_values.sum()} total missing values. Dropping rows with NaNs.\")\n", " df.dropna(subset=['Target', 'proteina', 'Y', 'PDB'], inplace=True)\n", " print(f\"Dataset now has {len(df)} rows after cleaning.\")\n", " \n", " df['Target'] = df['Target'].astype(str)\n", " df['proteina'] = df['proteina'].astype(str)\n", " df['Y'] = pd.to_numeric(df['Y'], errors='coerce')\n", " df['PDB'] = df['PDB'].astype(str)\n", " df['Subgroup'] = df['Subgroup'].astype(str)\n", " df['SourceDataSet'] = df['SourceDataSet'].astype(str) # NEW\n", "\n", " df.dropna(subset=['Y'], inplace=True)\n", " \n", " print(f\"\\nāœ… Data is loaded and cleaned. Final dataset size: {len(df)} samples.\")\n", " df.info()\n", "\n", "except FileNotFoundError:\n", " print(\"āŒ ERROR: 'Data.csv' not found.\")\n", " df = pd.DataFrame(columns=['Target', 'proteina', 'Y', 'PDB'])\n", "except KeyError as e:\n", " print(f\"āŒ ERROR: Column {e} not found in 'Data.csv'.\")\n", " df = pd.DataFrame(columns=['Target', 'proteina', 'Y', 'PDB'])\n", "\n", "# ==============================================================================\n", "# 4. EXPERIMENT CONFIGURATION\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 4. CONFIGURING THE EXPERIMENT ...\")\n", "\n", "CV_STRATEGY = 'SequenceSimilarity' \n", "N_SPLITS = 5\n", "EPOCHS = 30\n", "PATIENCE = 15\n", "BATCH_SIZE = 1\n", "LEARNING_RATE = 1e-4\n", "PROJECTED_SIZE = 256\n", "DROPOUT = 0.1\n", "\n", "LORA_RANK = 8\n", "LORA_ALPHA = 16\n", "LORA_DROPOUT = 0.1\n", "USE_LORA_FOR_TRAINING = True\n", "\n", "# Sequence Similarity CV Configs\n", "SIMILARITY_THRESHOLD = 0.3\n", "K_MER_SIZE = 3\n", "CLUSTERING_METHOD = 'agglomerative'\n", "\n", "RESULTS_DIR = \"cv_results_lora_seqsim\" # Use a unique name\n", "os.makedirs(RESULTS_DIR, exist_ok=True)\n", "print(f\"Results will be saved in: {RESULTS_DIR}\")\n", "\n", "print(f\"āœ… Configuration set. Using CV Strategy: {CV_STRATEGY}\")\n", "print(f\" Epochs: {EPOCHS}, Patience: {PATIENCE}\")\n", "if USE_LORA_FOR_TRAINING:\n", " print(f\" LoRA Enabled: Rank={LORA_RANK}, Alpha={LORA_ALPHA}, Dropout={LORA_DROPOUT}\")\n", "else:\n", " print(f\" LoRA Disabled: Only projection head will be trained.\")\n", "\n", "# ==============================================================================\n", "# 5. ADAPTED MODEL AND DATASET\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 5. DEFINING ADAPTED MODEL AND DATASET ...\")\n", "\n", "class BALMProjectionHead(nn.Module):\n", " def __init__(self, embedding_size, projected_size, projected_dropout):\n", " super().__init__()\n", " self.protein_projection = nn.Linear(embedding_size, projected_size)\n", " self.proteina_projection = nn.Linear(embedding_size, projected_size)\n", " self.dropout = nn.Dropout(projected_dropout)\n", " self.loss_fn = nn.MSELoss()\n", " print(f\"āœ… BALMProjectionHead initialized with embedding_size: {embedding_size}\")\n", "\n", " def forward(self, protein_embedding, proteina_embedding, labels=None):\n", " protein_embedding = self.dropout(protein_embedding)\n", " proteina_embedding = self.dropout(proteina_embedding)\n", " protein_projected = self.protein_projection(protein_embedding)\n", " proteina_projected = self.proteina_projection(proteina_embedding)\n", " protein_projected = F.normalize(protein_projected, p=2, dim=1)\n", " proteina_projected = F.normalize(proteina_projected, p=2, dim=1)\n", " cosine_similarity = F.cosine_similarity(protein_projected, proteina_projected)\n", " cosine_similarity = torch.clamp(cosine_similarity, -0.9999, 0.9999)\n", " output = {\"cosine_similarity\": cosine_similarity}\n", " if labels is not None:\n", " loss = self.loss_fn(cosine_similarity, labels)\n", " output[\"loss\"] = loss\n", " return output\n", "\n", "class BALMForLoRAFinetuning(nn.Module):\n", " def __init__(self, esm_model, esm_tokenizer, projected_size, projected_dropout, pkd_bounds):\n", " super().__init__()\n", " self.esm_model = esm_model\n", " self.esm_tokenizer = esm_tokenizer\n", " self.embedding_size = self.esm_model.config.hidden_size\n", " self.projection_head = BALMProjectionHead(self.embedding_size, projected_size, projected_dropout)\n", " self.pkd_lower, self.pkd_upper = pkd_bounds\n", " self.pkd_range = self.pkd_upper - self.pkd_lower\n", " print(f\"āœ… BALMForLoRAFinetuning model initialized.\")\n", "\n", " def _get_esm_embeddings(self, sequences: List[str]):\n", " cls_token = self.esm_tokenizer.cls_token\n", " processed_seqs = [s.replace('|', f\"{cls_token}{cls_token}\") for s in sequences]\n", " \n", " inputs = self.esm_tokenizer(processed_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024)\n", " inputs = {k: v.to(self.esm_model.device) for k, v in inputs.items()}\n", " outputs = self.esm_model(**inputs)\n", " \n", " # Pooling logic (remains Mean Pooling, which now covers the internal CLS tokens)\n", " hidden_states = outputs.last_hidden_state\n", " mask = inputs['attention_mask'].unsqueeze(-1).expand(hidden_states.size()).float()\n", " return torch.sum(hidden_states * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)\n", "\n", " def forward(self, batch_input):\n", " protein_emb = self._get_esm_embeddings(batch_input[\"protein_sequence\"])\n", " proteina_emb = self._get_esm_embeddings(batch_input[\"proteina_sequence\"])\n", " labels = batch_input.get(\"labels\")\n", " proj_output = self.projection_head(protein_emb, proteina_emb, labels)\n", " return {\n", " \"cosine_similarity\": proj_output[\"cosine_similarity\"],\n", " \"original_pkds\": batch_input[\"original_pkds\"],\n", " \"pdb_groups\": batch_input[\"pdb_groups\"],\n", " \"subgroups\": batch_input[\"subgroups\"],\n", " \"source_dataset\": batch_input[\"source_dataset\"], # NEW\n", " \"loss\": proj_output.get(\"loss\")\n", " }\n", " if \"loss\" in proj_output:\n", "\n", " output[\"loss\"] = proj_output[\"loss\"]\n", " return output\n", "\n", "class ProteinPairSequenceDataset(Dataset):\n", " def __init__(self, dataframe, pkd_bounds):\n", " self.data = dataframe.reset_index(drop=True)\n", " self.pkd_lower, self.pkd_upper = pkd_bounds\n", " self.pkd_range = self.pkd_upper - self.pkd_lower\n", "\n", " def __len__(self):\n", " return len(self.data)\n", "\n", " def __getitem__(self, idx):\n", " row = self.data.iloc[idx]\n", " original_pkd = float(row['Y'])\n", " scaled_label = ((original_pkd - self.pkd_lower) / self.pkd_range) * 2 - 1\n", " return {\n", " \"protein_sequence\": row[\"Target\"],\n", " \"proteina_sequence\": row[\"proteina\"],\n", " \"labels\": torch.tensor(scaled_label, dtype=torch.float32),\n", " \"original_pkds\": torch.tensor(original_pkd, dtype=torch.float32),\n", " \"pdb_groups\": row[\"PDB\"],\n", " \"subgroups\": row[\"Subgroup\"],\n", " \"source_dataset\": row[\"SourceDataSet\"] # NEW\n", " }\n", "\n", "def collate_fn_sequences(batch):\n", " return {\n", " \"protein_sequence\": [item['protein_sequence'] for item in batch],\n", " \"proteina_sequence\": [item['proteina_sequence'] for item in batch],\n", " \"labels\": torch.stack([item['labels'] for item in batch]),\n", " \"original_pkds\": torch.stack([item['original_pkds'] for item in batch]),\n", " \"pdb_groups\": [item['pdb_groups'] for item in batch], # NEW\n", " \"subgroups\": [item['subgroups'] for item in batch], # NEW\n", " \"source_dataset\": [item['source_dataset'] for item in batch] # NEW\n", " }\n", "\n", "# ==============================================================================\n", "# 6. METRICS, EVALUATION & SEQUENCE SIMILARITY HELPERS (NEW)\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 6. DEFINING METRICS, EVALUATION & SEQUENCE HELPERS ...\")\n", "\n", "def concordance_index(y_true, y_pred):\n", " try:\n", " y_true, y_pred = np.array(y_true), np.array(y_pred)\n", " n = len(y_true)\n", " if n < 2: return 1.0\n", " concordant, discordant, tied_pred = 0, 0, 0\n", " for i in range(n):\n", " for j in range(i + 1, n):\n", " if y_true[i] == y_true[j]: continue\n", " if y_pred[i] == y_pred[j]:\n", " tied_pred += 1\n", " continue\n", " if (y_pred[i] > y_pred[j] and y_true[i] > y_true[j]) or \\\n", " (y_pred[i] < y_pred[j] and y_true[i] < y_true[j]):\n", " concordant += 1\n", " else:\n", " discordant += 1\n", " total_pairs = concordant + discordant + tied_pred\n", " return 1.0 if total_pairs == 0 else (concordant + 0.5 * tied_pred) / total_pairs\n", " except Exception:\n", " return np.nan\n", "\n", "def evaluate_model(model, data_loader, device, pkd_bounds):\n", " model.eval()\n", " all_labels, all_preds_pkd, all_pdb_groups, all_subgroups, all_sources = [], [], [], [], [] \n", " pkd_lower, pkd_range = pkd_bounds[0], pkd_bounds[1] - pkd_bounds[0]\n", " \n", " with torch.no_grad():\n", " for batch in data_loader:\n", " batch_gpu = {k: v for k, v in batch.items()}\n", " for key in [\"labels\", \"original_pkds\"]:\n", " if key in batch_gpu: batch_gpu[key] = batch_gpu[key].to(device)\n", " outputs = model(batch_gpu)\n", " cosine_sim = outputs['cosine_similarity'].cpu()\n", " preds_pkd = ((cosine_sim + 1) / 2) * pkd_range + pkd_lower\n", " all_labels.extend(batch['original_pkds'].numpy())\n", " all_preds_pkd.extend(preds_pkd.numpy())\n", " all_pdb_groups.extend(batch['pdb_groups']) \n", " all_subgroups.extend(batch['subgroups']) \n", " all_sources.extend(batch['source_dataset']) # NEW\n", "\n", " labels, preds = np.array(all_labels), np.array(all_preds_pkd)\n", " pdb_groups_arr = np.array(all_pdb_groups)\n", " subgroups_arr = np.array(all_subgroups)\n", " sources_arr = np.array(all_sources) # NEW\n", " \n", " mask = ~np.isnan(labels) & ~np.isnan(preds)\n", " labels, preds = labels[mask], preds[mask]\n", " \n", " if len(labels) < 2:\n", " return {\"rmse\": np.nan, \"pearson\": np.nan, \"spearman\": np.nan, \"ci\": np.nan}, labels, preds, pdb_groups_arr, subgroups_arr, sources_arr\n", " return {\n", " 'rmse': np.sqrt(mean_squared_error(labels, preds)),\n", " 'pearson': pearsonr(labels, preds)[0],\n", " 'spearman': spearmanr(labels, preds)[0],\n", " 'ci': concordance_index(labels, preds)\n", " }, labels, preds, pdb_groups_arr, subgroups_arr, sources_arr\n", "\n", "def plot_regression(y_true, y_pred, metrics, title, filename=None):\n", " plt.figure(figsize=(8, 8))\n", " sns.scatterplot(x=y_true, y=y_pred, alpha=0.6, edgecolor='k', s=50)\n", " min_val, max_val = min(np.min(y_true), np.min(y_pred)), max(np.max(y_true), np.max(y_pred))\n", " plt.plot([min_val, max_val], [min_val, max_val], 'r--', lw=2, label='Identity Line (y=x)')\n", " metrics_text = f\"RMSE: {metrics['rmse']:.4f}\\nPearson: {metrics['pearson']:.4f}\\nSpearman: {metrics['spearman']:.4f}\\nCI: {metrics['ci']:.4f}\"\n", " plt.text(0.05, 0.95, metrics_text, transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.5))\n", " plt.title(title, fontsize=16)\n", " plt.xlabel(\"Actual pKd\", fontsize=14)\n", " plt.ylabel(\"Predicted pKd\", fontsize=14)\n", " plt.legend()\n", " plt.grid(True)\n", " plt.tight_layout()\n", " if filename: # NEW\n", " plt.savefig(filename) # NEW\n", " print(f\"Plot saved to {filename}\") # NEW\n", " plt.show()\n", "\n", "def compute_kmer_similarity_matrix(sequences, k=3):\n", " \"\"\"\n", " Compute pairwise sequence similarity matrix using k-mer based Jaccard similarity.\n", " This is a biologically relevant method for sequence similarity computation.\n", " \"\"\"\n", " print(f\"Computing {k}-mer Jaccard similarity matrix for {len(sequences)} sequences...\")\n", " \n", " def get_kmers(sequence, k_size):\n", " \"\"\"Extract k-mers from a sequence.\"\"\"\n", " if len(sequence) < k_size:\n", " return {sequence}\n", " return {sequence[i:i+k_size] for i in range(len(sequence) - k_size + 1)}\n", " \n", " def jaccard_similarity(seq1, seq2):\n", " \"\"\"Compute Jaccard similarity between two sequences based on k-mers.\"\"\"\n", " kmers1 = get_kmers(seq1, k)\n", " kmers2 = get_kmers(seq2, k)\n", " if not kmers1 or not kmers2:\n", " return 0.0\n", " intersection = len(kmers1 & kmers2)\n", " union = len(kmers1 | kmers2)\n", " return intersection / union if union > 0 else 0.0\n", " \n", " n_seqs = len(sequences)\n", " similarity_matrix = np.zeros((n_seqs, n_seqs))\n", " \n", " # Compute similarity matrix with progress bar\n", " for i in tqdm(range(n_seqs), desc=\"Computing similarities\"):\n", " for j in range(i, n_seqs):\n", " if i == j:\n", " similarity_matrix[i, j] = 1.0\n", " else:\n", " sim = jaccard_similarity(sequences[i], sequences[j])\n", " similarity_matrix[i, j] = sim\n", " similarity_matrix[j, i] = sim\n", " \n", " avg_similarity = np.mean(similarity_matrix[np.triu_indices_from(similarity_matrix, k=1)])\n", " print(f\"Average pairwise similarity: {avg_similarity:.3f}\")\n", " \n", " return similarity_matrix\n", "\n", "def cluster_sequences_by_similarity(similarity_matrix, threshold=0.3, method='agglomerative'):\n", " \"\"\"\n", " Cluster sequences based on similarity matrix using specified clustering method.\n", " Uses distance threshold derived from similarity threshold.\n", " \"\"\"\n", " print(f\"Clustering sequences using {method} method with similarity threshold {threshold}\")\n", " \n", " # Convert similarity to distance matrix\n", " distance_matrix = 1.0 - similarity_matrix\n", " \n", " try:\n", " if method == 'agglomerative':\n", " # Use distance threshold - sequences with similarity > threshold will be in same cluster\n", " distance_threshold = 1.0 - threshold\n", " \n", " print(f\"Using distance threshold: {distance_threshold}\")\n", " print(f\"Distance matrix stats - Min: {distance_matrix.min():.3f}, Max: {distance_matrix.max():.3f}, Mean: {distance_matrix.mean():.3f}\")\n", " \n", " clustering = AgglomerativeClustering(\n", " n_clusters=None,\n", " distance_threshold=distance_threshold,\n", " linkage='average', # Average linkage works well for biological sequences\n", " metric='precomputed'\n", " )\n", " \n", " clusters = clustering.fit_predict(distance_matrix)\n", " print(f\"AgglomerativeClustering completed successfully\")\n", " \n", " else:\n", " raise ValueError(f\"Unsupported clustering method: {method}\")\n", " \n", " if clusters is None:\n", " raise ValueError(\"Clustering failed - returned None\")\n", " \n", " n_clusters = len(set(clusters))\n", " print(f\"Created {n_clusters} clusters\")\n", " \n", " cluster_sizes = [np.sum(clusters == i) for i in set(clusters)]\n", " print(f\"Cluster size distribution - Min: {min(cluster_sizes)}, Max: {max(cluster_sizes)}, Mean: {np.mean(cluster_sizes):.1f}\")\n", " \n", " return clusters\n", " \n", " except Exception as e:\n", " print(f\"āŒ Clustering failed with error: {e}\")\n", " print(\"šŸ”„ Falling back to individual clusters (each sequence in its own cluster)\")\n", " \n", " n_sequences = similarity_matrix.shape[0]\n", " fallback_clusters = np.arange(n_sequences)\n", " \n", " print(f\"Created {n_sequences} individual clusters (fallback)\")\n", " return fallback_clusters\n", "\n", "# ==============================================================================\n", "# 7. CROSS-VALIDATION TRAINING (MODIFIED FOR 80/20 SPLIT + TRAIN LOSS PATIENCE)\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(f\"... 7. RUNNING '{CV_STRATEGY}' STRATEGY WITH LoRA (80/20 Split with Training Loss Patience) ...\")\n", "\n", "def run_sequence_similarity_cv_with_lora():\n", " if df.empty:\n", " print(\"āš ļø Skipping training due to empty dataframe.\")\n", " return\n", "\n", " device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " PKD_LOWER_BOUND, PKD_UPPER_BOUND = df['Y'].min(), df['Y'].max()\n", " PKD_BOUNDS = (PKD_LOWER_BOUND, PKD_UPPER_BOUND)\n", " print(f\"Global pKd bounds for scaling: [{PKD_LOWER_BOUND:.2f}, {PKD_UPPER_BOUND:.2f}]\")\n", "\n", " # --- Steps 1-4: Data Preparation and Clustering ---\n", " print(\"\\nšŸ§¬ Step 1: Extracting unique protein sequences...\")\n", " sequence_columns = ['Target', 'proteina']\n", " all_sequences = set()\n", " sequence_to_records = {}\n", " for idx, row in df.iterrows():\n", " for col in sequence_columns:\n", " seq = str(row[col]).strip()\n", " if seq:\n", " all_sequences.add(seq)\n", " if seq not in sequence_to_records:\n", " sequence_to_records[seq] = []\n", " sequence_to_records[seq].append(idx)\n", " unique_sequences = sorted(list(all_sequences))\n", " print(f\"Found {len(unique_sequences)} unique protein sequences.\")\n", "\n", " print(\"\\nšŸ§¬ Step 2: Computing sequence similarity matrix...\")\n", " similarity_matrix = compute_kmer_similarity_matrix(unique_sequences, k=K_MER_SIZE)\n", "\n", " print(\"\\nšŸ§¬ Step 3: Clustering sequences...\")\n", " clusters = cluster_sequences_by_similarity(similarity_matrix, threshold=SIMILARITY_THRESHOLD, method=CLUSTERING_METHOD)\n", " \n", " cluster_to_records = {}\n", " for seq_idx, cluster_id in enumerate(clusters):\n", " seq = unique_sequences[seq_idx]\n", " if cluster_id not in cluster_to_records:\n", " cluster_to_records[cluster_id] = []\n", " cluster_to_records[cluster_id].extend(sequence_to_records[seq])\n", " \n", " for cluster_id in cluster_to_records:\n", " cluster_to_records[cluster_id] = list(set(cluster_to_records[cluster_id]))\n", " # --- End of omitted code ---\n", "\n", "\n", " # --- Step 5: Perform Cluster-based Cross-Validation ---\n", " print(f\"\\nšŸš€ Starting {N_SPLITS}-Fold **{CV_STRATEGY}** Cross-Validation with LoRA fine-tuning...\")\n", " \n", " # NOTE: Model initialization moved INSIDE the loop to prevent weight leakage.\n", " \n", " cluster_ids = list(cluster_to_records.keys())\n", " kf = KFold(n_splits=N_SPLITS, shuffle=True, random_state=42)\n", " \n", " fold_results, all_folds_y_true, all_folds_y_pred = [], [], []\n", " \n", " for fold, (train_cluster_idx, test_cluster_idx) in enumerate(kf.split(cluster_ids)):\n", " fold_num = fold + 1\n", " print(f\"\\n===== STRATEGY: {CV_STRATEGY} | FOLD {fold_num}/{N_SPLITS} =====\")\n", " \n", " # --- CRITICAL FIX: RE-INITIALIZE MODEL HERE ---\n", " # This ensures the LoRA weights are reset for every fold.\n", " print(f\"šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold {fold_num}...\")\n", " extractor = ProteinEmbeddingExtractor(\n", " device=device, use_lora=USE_LORA_FOR_TRAINING, lora_rank=LORA_RANK,\n", " lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT\n", " )\n", " esm_model, esm_tokenizer = extractor.get_model_and_tokenizer()\n", " # ----------------------------------------------\n", "\n", " fold_model_path = os.path.join(RESULTS_DIR, f\"best_model_fold_{fold_num}.pth\")\n", " \n", " train_clusters = [cluster_ids[i] for i in train_cluster_idx]\n", " test_clusters = [cluster_ids[i] for i in test_cluster_idx]\n", "\n", " # Create a map from each unique sequence to its group ('train' or 'test')\n", " sequence_to_group = {}\n", " train_cluster_set = set(train_clusters)\n", " test_cluster_set = set(test_clusters)\n", "\n", " for seq_idx, cluster_id in enumerate(clusters):\n", " sequence = unique_sequences[seq_idx]\n", " if cluster_id in train_cluster_set:\n", " sequence_to_group[sequence] = 'train'\n", " elif cluster_id in test_cluster_set:\n", " sequence_to_group[sequence] = 'test'\n", "\n", " # Iterate through the main dataframe and assign rows\n", " train_indices, test_indices = [], []\n", " \n", " for idx, row in df.iterrows():\n", " seq_target = str(row['Target']).strip()\n", " seq_proteina = str(row['proteina']).strip()\n", "\n", " cluster_target_id = sequence_to_group.get(seq_target)\n", " cluster_proteina_id = sequence_to_group.get(seq_proteina)\n", "\n", " # Strict Leakage Rule: If EITHER protein is in a test cluster, it goes to test.\n", " if cluster_target_id == 'test' or cluster_proteina_id == 'test':\n", " test_indices.append(idx)\n", " else:\n", " train_indices.append(idx)\n", "\n", " # Create dataframes\n", " train_df = df.loc[train_indices]\n", " test_df = df.loc[test_indices]\n", " \n", " assert len(set(train_indices).intersection(set(test_indices))) == 0, \"Overlap found between train and test sets!\"\n", "\n", " print(f\" Total samples allocated: {len(train_df) + len(test_df)}\")\n", " print(f\" Train: {len(train_df)} samples\")\n", " print(f\" Test: {len(test_df)} samples\")\n", "\n", " # --- FIX: Calculate Scaling Bounds using ONLY Training Data (Prevents Data Sniffing) ---\n", " fold_pkd_min = train_df['Y'].min()\n", " fold_pkd_max = train_df['Y'].max()\n", " FOLD_PKD_BOUNDS = (fold_pkd_min, fold_pkd_max)\n", " print(f\" Fold Scaling Bounds (derived from train set): {FOLD_PKD_BOUNDS}\")\n", "\n", " # --- DataLoaders ---\n", " train_dataset = ProteinPairSequenceDataset(train_df, FOLD_PKD_BOUNDS)\n", " test_dataset = ProteinPairSequenceDataset(test_df, FOLD_PKD_BOUNDS)\n", " train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn_sequences)\n", " test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE*2, shuffle=False, collate_fn=collate_fn_sequences)\n", "\n", " # --- Model, Optimizer ---\n", " model = BALMForLoRAFinetuning(\n", " esm_model=esm_model, esm_tokenizer=esm_tokenizer,\n", " projected_size=PROJECTED_SIZE, projected_dropout=DROPOUT,\n", " pkd_bounds=FOLD_PKD_BOUNDS\n", " ).to(device)\n", " \n", " optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)\n", " print(f\" Optimizing {sum(p.numel() for p in model.parameters() if p.requires_grad)} trainable parameters.\")\n", "\n", " # --- Training Loop ---\n", " best_train_loss = float('inf')\n", " patience_counter = 0\n", "\n", " for epoch in range(EPOCHS):\n", " model.train()\n", " total_train_loss = 0.0\n", " progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS}\", leave=False)\n", " \n", " for batch in progress_bar:\n", " optimizer.zero_grad()\n", " batch_gpu = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}\n", " outputs = model(batch_gpu)\n", " loss = outputs['loss']\n", " loss.backward()\n", " optimizer.step()\n", " total_train_loss += loss.item()\n", " progress_bar.set_postfix({'batch_loss': f\"{loss.item():.4f}\"})\n", " \n", " avg_train_loss = total_train_loss / len(train_loader)\n", " # Evaluate on test set (for monitoring only, not for gradient updates)\n", " val_metrics, _, _, _, _, _ = evaluate_model(model, test_loader, device, FOLD_PKD_BOUNDS)\n", " print(f\"Epoch {epoch+1}/{EPOCHS} | Train Loss: {avg_train_loss:.4f} | Test RMSE: {val_metrics['rmse']:.4f}\")\n", " \n", " # Save best model based on Training Loss (to avoid validation leakage in parameter selection)\n", " if avg_train_loss < best_train_loss:\n", " best_train_loss = avg_train_loss\n", " patience_counter = 0\n", " torch.save(model.state_dict(), fold_model_path)\n", " # print(f\" -> New best model saved with Train Loss: {best_train_loss:.4f}\")\n", " else:\n", " patience_counter += 1\n", " \n", " if patience_counter >= PATIENCE:\n", " print(f\"ā³ Early stopping triggered based on training loss at epoch {epoch+1}.\")\n", " break\n", " \n", " print(f\"Fold {fold_num} training complete.\") \n", " # --- Final Evaluation on Test Set ---\n", " print(\"Loading best model (based on lowest train loss) for final TEST evaluation...\")\n", " model.load_state_dict(torch.load(fold_model_path)) \n", " \n", " final_test_metrics, y_true_fold, y_pred_fold, pdb_groups_fold, subgroups_fold, sources_fold = evaluate_model(model, test_loader, device, PKD_BOUNDS)\n", " print(f\"Fold {fold_num} Final Test Metrics:\")\n", " for key, value in final_test_metrics.items():\n", " print(f\" - {key.upper()}: {value:.4f}\")\n", "\n", " fold_results.append(final_test_metrics)\n", " all_folds_y_true.extend(y_true_fold)\n", " all_folds_y_pred.extend(y_pred_fold)\n", "\n", " # --- Save Predictions for the Fold ---\n", " # MODIFIED: Create a more detailed DataFrame\n", " pred_df = pd.DataFrame({\n", " 'label': y_true_fold, \n", " 'prediction': y_pred_fold, \n", " 'residual': y_true_fold - y_pred_fold,\n", " 'abs_residual': np.abs(y_true_fold - y_pred_fold),\n", " 'PDB': pdb_groups_fold, \n", " 'Subgroup': subgroups_fold,\n", " 'Source Data Set': sources_fold # Preserved column\n", " })\n", " # MODIFIED: Save to the RESULTS_DIR\n", " pred_filename = os.path.join(RESULTS_DIR, f\"fold_{fold_num}_predictions.csv\")\n", " pred_df.to_csv(pred_filename, index=False)\n", " print(f\" šŸ’¾ Predictions for fold {fold_num} saved to '{pred_filename}'\")\n", "\n", "\n", " # --- Aggregate and Display Final Results ---\n", " print(f\"\\nšŸ“ˆ Generating combined regression plot for {CV_STRATEGY}\")\n", " combined_y_true = np.array(all_folds_y_true)\n", " combined_y_pred = np.array(all_folds_y_pred)\n", "\n", " overall_metrics = {\n", " 'rmse': np.sqrt(mean_squared_error(combined_y_true, combined_y_pred)),\n", " 'pearson': pearsonr(combined_y_true, combined_y_pred)[0],\n", " 'spearman': spearmanr(combined_y_true, combined_y_pred)[0],\n", " 'ci': concordance_index(combined_y_true, combined_y_pred)\n", " }\n", " \n", " plot_regression(\n", " combined_y_true,\n", " combined_y_pred,\n", " overall_metrics,\n", " title=f\"Overall Regression for {CV_STRATEGY} ({N_SPLITS}-Fold CV) with LoRA Fine-tuning\",\n", " filename=os.path.join(RESULTS_DIR, \"overall_regression_plot.png\")\n", " )\n", "\n", " print(f\"\\nāœ… Final Results Summary for {CV_STRATEGY} with LoRA Fine-tuning:\")\n", " results_df = pd.DataFrame(fold_results)\n", " summary_df = pd.DataFrame({'Mean': results_df.mean(), 'Std Dev': results_df.std()})\n", " print(summary_df)\n", "\n", " # NEW: Save the summary metrics to a CSV file\n", " summary_filename = os.path.join(RESULTS_DIR, \"cv_summary_metrics.csv\")\n", " summary_df.to_csv(summary_filename)\n", " print(f\"Overall CV summary metrics saved to {summary_filename}\")\n", " print(f\"Fold model checkpoints and prediction files are saved in the '{RESULTS_DIR}' directory.\")\n", "\n", "# --- Call the new CV function ---\n", "run_sequence_similarity_cv_with_lora()\n", "\n", "\n", "# ==============================================================================\n", "# 8. FINAL NOTES AND RECOMMENDATIONS\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 8. FINAL NOTES AND RECOMMENDATIONS ...\")\n", "print(\"\"\"\n", "šŸŽ‰ CONGRATULATIONS! You've successfully implemented BALM with on-the-fly ESM-2 embeddings\n", "and **Parameter-Efficient Fine-Tuning (PEFT) using LoRA**!\n", "\n", "šŸ“Š KEY ADVANTAGES OF THIS APPROACH:\n", "ā€¢ **Efficiency:** Fine-tunes only a tiny fraction of the ESM-2 model parameters (<1%).\n", "ā€¢ **Speed & Memory:** Significantly faster and less memory-intensive than full fine-tuning.\n", "ā€¢ **Performance:** Adapts the powerful ESM-2 model to your specific data, often leading to better performance.\n", "ā€¢ **No Static Cache:** By computing embeddings dynamically, the model fully leverages the fine-tuned LoRA adapters during training.\n", "\n", "šŸš€ PERFORMANCE OPTIMIZATION TIPS:\n", "1. **LoRA Hyperparameters:**\n", " * `lora_rank (r)`: Controls the capacity of the adapters. Higher `r` means more parameters. Typical values are 8-64.\n", " * `lora_alpha`: Acts as a scaling factor. A common heuristic is to set `lora_alpha` to be twice the `lora_rank`. [7]\n", " * `target_modules`: Experiment with adding other modules like `dense` to the target list in the `LoraConfig`. [1]\n", "2. **Learning Rate:** LoRA can sometimes handle slightly higher learning rates than full fine-tuning. Experiment with values from 5e-5 to 5e-4.\n", "3. **Batch Size & Gradient Accumulation:** If your GPU memory is a constraint, you can simulate a larger batch size by using gradient accumulation.\n", "\n", "šŸ’¾ SAVING AND LOADING THE FINE-TUNED MODEL:\n", "ā€¢ After training, you can save the trained LoRA adapters separately from the base model. This is very efficient for storage. [3]\n", "ā€¢ To save: `model.esm_model.save_pretrained(\"path/to/my_lora_adapters\")` [3]\n", "ā€¢ To load: First, load the base ESM-2 model. Then, use `PeftModel.from_pretrained(base_model, \"path/to/my_lora_adapters\")`.\n", "\n", "šŸŽÆ NEXT STEPS:\n", "1. **Hyperparameter Search:** Systematically tune LoRA parameters (`r`, `lora_alpha`) and the learning rate for optimal performance.\n", "2. **Experiment with Larger ESM-2 Models:** LoRA makes fine-tuning even the largest ESM-2 models (e.g., 3B parameters) feasible on a single high-end GPU.\n", "3. **Advanced PEFT Techniques:** Explore other PEFT methods like IA3 or full parameter-efficient fine-tuning for comparison.\n", "\"\"\")\n", "print(\"\\nāœ… LoRA PEFT implementation complete! Happy parameter-efficient protein modeling! šŸ§¬šŸ¤–\")" ] }, { "cell_type": "code", "execution_count": 2, "id": "da181303", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "================================================================================\n", "... 0. PEFT LIBRARY INITIALIZED ...\n", " Ready for Parameter-Efficient Fine-Tuning with LoRA!\n", "šŸ§¬ BALM Adaptation for ESM-2 with LoRA Fine-Tuning\n", "āš” Fast Training with 5-Fold Cross-Validation\n", "================================================================================\n", "šŸ”’ Setting up reproducibility with seed 42\n", "āœ… Reproducibility setup complete\n", "\n", "================================================================================\n", "... 3. LOADING DATA FROM Data.csv ...\n", "Successfully loaded Data.csv with 12062 rows.\n", "\n", "--- Initial Data Check ---\n", " Target \\\n", "0 PKFTKCRSPERETFSCHWTLGPIQLFYTRRNTQEWTQEWKECPDYV... \n", "1 QDNSRYTHFLTQHYDAKPQGRDDRYCESIMRRRGLTSPCKDINTFI... \n", "2 KSFPEVVGKTVDQAREYFTLHYPQYDVYFLPEGSPVTLDLRYNRVR... \n", "3 TNTVAAYNLTWKSTNFKTILEWEPKPVNQVYTVQISTKSGDWKSKC... \n", "4 PIVQNLQGQMVHQAISPRTLNAWVKVVEEKAFSPEVIPMFSALSEG... \n", "\n", " proteina Y PDB \\\n", "0 FPTIPLSRLFDNAMLRAHRLHQLAFDTYQEFEEAYIPKEQKYSFLQ... 9.045757 1A22 \n", "1 SLDIQSLDIQCEELSDARWAELLPLLQQCQVVRLDDCGLTEARCKD... 15.301030 1A4Y \n", "2 CGVPAIQPVLSGLIVNGEEAVPGSWPWQVSLQDKTGFHFCGGSLIN... 11.826814 1ACB \n", "3 DIKMTQSPSSMYASLGERVTITCKASQDIRKYLNWYQQKPWKSPKT... 8.468521 1AHW \n", "4 MVNPTVFFDIAVDGEPLGRVSFELFADKVPKTAENFRALSTGEKGF... 4.920819 1AK4 \n", "\n", " Subgroup SourceDataSet \n", "0 others SKEMPI v2.0 \n", "1 others SKEMPI v2.0 \n", "2 others SKEMPI v2.0 \n", "3 others SKEMPI v2.0 \n", "4 others SKEMPI v2.0 \n", "\n", "Missing values per column:\n", " Target 39\n", "proteina 28\n", "Y 0\n", "PDB 0\n", "Subgroup 0\n", "SourceDataSet 0\n", "dtype: int64\n", "\n", "Warning: Found 67 total missing values. Dropping rows with NaNs.\n", "Dataset now has 12019 rows after cleaning.\n", "\n", "āœ… Data is loaded and cleaned. Final dataset size: 12019 samples.\n", "\n", "Index: 12019 entries, 0 to 12061\n", "Data columns (total 6 columns):\n", " # Column Non-Null Count Dtype \n", "--- ------ -------------- ----- \n", " 0 Target 12019 non-null object \n", " 1 proteina 12019 non-null object \n", " 2 Y 12019 non-null float64\n", " 3 PDB 12019 non-null object \n", " 4 Subgroup 12019 non-null object \n", " 5 SourceDataSet 12019 non-null object \n", "dtypes: float64(1), object(5)\n", "memory usage: 657.3+ KB\n", "\n", "================================================================================\n", "... 4. CONFIGURING THE EXPERIMENT ...\n", "Results will be saved in: cv_results_lora_seqsim\n", "āœ… Configuration set. Using CV Strategy: SequenceSimilarity\n", " Epochs: 30, Patience: 15\n", " LoRA Enabled: Rank=8, Alpha=16, Dropout=0.1\n", "\n", "================================================================================\n", "... 5. DEFINING ADAPTED MODEL AND DATASET ...\n", "\n", "================================================================================\n", "... 6. DEFINING METRICS, EVALUATION & SEQUENCE HELPERS ...\n", "\n", "================================================================================\n", "... 7. RUNNING 'SequenceSimilarity' STRATEGY WITH LoRA (80/20 Split with Training Loss Patience) ...\n", "Global pKd bounds for scaling: [1.32, 15.70]\n", "\n", "šŸ§¬ Step 1: Extracting unique protein sequences...\n", "Found 11598 unique protein sequences.\n", "\n", "šŸ§¬ Step 2: Computing sequence similarity matrix...\n", "Computing 3-mer Jaccard similarity matrix for 11598 sequences...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Computing similarities: 100%|ā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆā–ˆ| 11598/11598 [1:00:04<00:00, 3.22it/s] \n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Average pairwise similarity: 0.028\n", "\n", "šŸ§¬ Step 3: Clustering sequences...\n", "Clustering sequences using agglomerative method with similarity threshold 0.3\n", "Using distance threshold: 0.7\n", "Distance matrix stats - Min: 0.000, Max: 1.000, Mean: 0.972\n", "Hierarchical clustering completed successfully (deterministic)\n", "Created 2678 clusters\n", "Cluster size distribution - Min: 1, Max: 1060, Mean: 4.3\n", "\n", "šŸš€ Starting 5-Fold **SequenceSimilarity** Cross-Validation with LoRA fine-tuning...\n", "\n", "Created 2678 clusters\n", " Step 4: Creating 5-fold splits...\n", " Fold 1: Train=6516, Test=5503\n", " Fold 2: Train=7503, Test=4516\n", " Fold 3: Train=7340, Test=4679\n", " Fold 4: Train=7331, Test=4688\n", " Fold 5: Train=9397, Test=2622\n", "\n", "šŸ”„ Beginning actual training across 5 folds...\n", "\n", "===== STRATEGY: SequenceSimilarity | FOLD 1/5 =====\n", "šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold 1...\n", "šŸ”§ Initializing ProteinEmbeddingExtractor on cuda\n", "šŸ“„ Loading ESM-2 model: facebook/esm2_t33_650M_UR50D with dtype: torch.float16\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n", "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "āœØ Applying LoRA with rank=8, alpha=16, dropout=0.1\n", " āœ… LoRA model loaded and configured. Trainable parameters:\n", "trainable params: 2,027,520 || all params: 654,381,461 || trainable%: 0.3098\n", " āœ… Model loaded. Embedding size from config: 1280\n", " Total samples allocated: 12019\n", " Train: 6516 samples\n", " Test: 5503 samples\n", " Using Global Scaling Bounds: (1.3187587626244128, 15.69897000433602)\n", "āœ… BALMProjectionHead initialized with embedding_size: 1280\n", "āœ… BALMForLoRAFinetuning model initialized.\n", " Optimizing 2683392 trainable parameters.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 1/30 | Train Loss: 0.0428 | Test RMSE: 1.7014 | Test Pearson: 0.6292\n", " -> New best model saved with Val RMSE: 1.7014\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 2/30 | Train Loss: 0.0302 | Test RMSE: 1.8045 | Test Pearson: 0.6331\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 3/30 | Train Loss: 0.0253 | Test RMSE: 1.6909 | Test Pearson: 0.6508\n", " -> New best model saved with Val RMSE: 1.6909\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 4/30 | Train Loss: 0.0213 | Test RMSE: 1.7378 | Test Pearson: 0.6570\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 5/30 | Train Loss: 0.0187 | Test RMSE: 1.7107 | Test Pearson: 0.6377\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 6/30 | Train Loss: 0.0170 | Test RMSE: 1.8107 | Test Pearson: 0.6385\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 7/30 | Train Loss: 0.0156 | Test RMSE: 1.7539 | Test Pearson: 0.6507\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 8/30 | Train Loss: 0.0143 | Test RMSE: 1.7244 | Test Pearson: 0.6537\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 9/30 | Train Loss: 0.0135 | Test RMSE: 1.6885 | Test Pearson: 0.6550\n", " -> New best model saved with Val RMSE: 1.6885\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 10/30 | Train Loss: 0.0124 | Test RMSE: 1.7479 | Test Pearson: 0.6446\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 11/30 | Train Loss: 0.0121 | Test RMSE: 1.7309 | Test Pearson: 0.6522\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 12/30 | Train Loss: 0.0113 | Test RMSE: 1.7429 | Test Pearson: 0.6304\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 13/30 | Train Loss: 0.0108 | Test RMSE: 1.7986 | Test Pearson: 0.6161\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 14/30 | Train Loss: 0.0104 | Test RMSE: 1.8265 | Test Pearson: 0.6088\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 15/30 | Train Loss: 0.0102 | Test RMSE: 1.7594 | Test Pearson: 0.6410\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 16/30 | Train Loss: 0.0098 | Test RMSE: 1.7327 | Test Pearson: 0.6386\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 17/30 | Train Loss: 0.0093 | Test RMSE: 1.8332 | Test Pearson: 0.6202\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 18/30 | Train Loss: 0.0090 | Test RMSE: 1.7435 | Test Pearson: 0.6247\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 19/30 | Train Loss: 0.0088 | Test RMSE: 1.8083 | Test Pearson: 0.6308\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 20/30 | Train Loss: 0.0088 | Test RMSE: 1.7652 | Test Pearson: 0.6330\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 21/30 | Train Loss: 0.0085 | Test RMSE: 1.7683 | Test Pearson: 0.6317\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 22/30 | Train Loss: 0.0084 | Test RMSE: 1.7539 | Test Pearson: 0.6220\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 23/30 | Train Loss: 0.0085 | Test RMSE: 1.7776 | Test Pearson: 0.6337\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 24/30 | Train Loss: 0.0079 | Test RMSE: 1.7249 | Test Pearson: 0.6325\n", "ā³ Early stopping triggered at epoch 24. Val RMSE has not improved for 15 epochs.\n", "Fold 1 training complete.\n", "Loading best model (based on lowest train loss) for final TEST evaluation...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\hs494\\AppData\\Local\\Temp\\ipykernel_46900\\1832587606.py:744: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " model.load_state_dict(torch.load(fold_model_path))\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Fold 1 Final Test Metrics:\n", " - RMSE: 1.6885\n", " - PEARSON: 0.6550\n", " - SPEARMAN: 0.6609\n", " - CI: 0.7377\n", " šŸ’¾ Predictions for fold 1 saved to 'cv_results_lora_seqsim\\fold_1_predictions.csv'\n", "\n", "===== STRATEGY: SequenceSimilarity | FOLD 2/5 =====\n", "šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold 2...\n", "šŸ”§ Initializing ProteinEmbeddingExtractor on cuda\n", "šŸ“„ Loading ESM-2 model: facebook/esm2_t33_650M_UR50D with dtype: torch.float16\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n", "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "āœØ Applying LoRA with rank=8, alpha=16, dropout=0.1\n", " āœ… LoRA model loaded and configured. Trainable parameters:\n", "trainable params: 2,027,520 || all params: 654,381,461 || trainable%: 0.3098\n", " āœ… Model loaded. Embedding size from config: 1280\n", " Total samples allocated: 12019\n", " Train: 7503 samples\n", " Test: 4516 samples\n", " Using Global Scaling Bounds: (1.3187587626244128, 15.69897000433602)\n", "āœ… BALMProjectionHead initialized with embedding_size: 1280\n", "āœ… BALMForLoRAFinetuning model initialized.\n", " Optimizing 2683392 trainable parameters.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 1/30 | Train Loss: 0.0441 | Test RMSE: 1.7670 | Test Pearson: 0.4830\n", " -> New best model saved with Val RMSE: 1.7670\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 2/30 | Train Loss: 0.0318 | Test RMSE: 1.6988 | Test Pearson: 0.5175\n", " -> New best model saved with Val RMSE: 1.6988\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 3/30 | Train Loss: 0.0245 | Test RMSE: 1.8133 | Test Pearson: 0.4926\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 4/30 | Train Loss: 0.0209 | Test RMSE: 1.7363 | Test Pearson: 0.5227\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 5/30 | Train Loss: 0.0183 | Test RMSE: 1.8534 | Test Pearson: 0.4780\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 6/30 | Train Loss: 0.0165 | Test RMSE: 1.8327 | Test Pearson: 0.5160\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 7/30 | Train Loss: 0.0158 | Test RMSE: 1.7264 | Test Pearson: 0.5440\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 8/30 | Train Loss: 0.0145 | Test RMSE: 1.8015 | Test Pearson: 0.5091\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 9/30 | Train Loss: 0.0137 | Test RMSE: 1.7813 | Test Pearson: 0.5071\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 10/30 | Train Loss: 0.0128 | Test RMSE: 1.7957 | Test Pearson: 0.5295\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 11/30 | Train Loss: 0.0119 | Test RMSE: 1.7887 | Test Pearson: 0.4990\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 12/30 | Train Loss: 0.0115 | Test RMSE: 1.7319 | Test Pearson: 0.5139\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 13/30 | Train Loss: 0.0109 | Test RMSE: 1.7580 | Test Pearson: 0.4985\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 14/30 | Train Loss: 0.0106 | Test RMSE: 1.8255 | Test Pearson: 0.4744\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 15/30 | Train Loss: 0.0101 | Test RMSE: 1.7823 | Test Pearson: 0.4961\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 16/30 | Train Loss: 0.0095 | Test RMSE: 1.7827 | Test Pearson: 0.4969\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 17/30 | Train Loss: 0.0091 | Test RMSE: 1.8194 | Test Pearson: 0.4644\n", "ā³ Early stopping triggered at epoch 17. Val RMSE has not improved for 15 epochs.\n", "Fold 2 training complete.\n", "Loading best model (based on lowest train loss) for final TEST evaluation...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\hs494\\AppData\\Local\\Temp\\ipykernel_46900\\1832587606.py:744: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " model.load_state_dict(torch.load(fold_model_path))\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Fold 2 Final Test Metrics:\n", " - RMSE: 1.6988\n", " - PEARSON: 0.5175\n", " - SPEARMAN: 0.4978\n", " - CI: 0.6721\n", " šŸ’¾ Predictions for fold 2 saved to 'cv_results_lora_seqsim\\fold_2_predictions.csv'\n", "\n", "===== STRATEGY: SequenceSimilarity | FOLD 3/5 =====\n", "šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold 3...\n", "šŸ”§ Initializing ProteinEmbeddingExtractor on cuda\n", "šŸ“„ Loading ESM-2 model: facebook/esm2_t33_650M_UR50D with dtype: torch.float16\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n", "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "āœØ Applying LoRA with rank=8, alpha=16, dropout=0.1\n", " āœ… LoRA model loaded and configured. Trainable parameters:\n", "trainable params: 2,027,520 || all params: 654,381,461 || trainable%: 0.3098\n", " āœ… Model loaded. Embedding size from config: 1280\n", " Total samples allocated: 12019\n", " Train: 7340 samples\n", " Test: 4679 samples\n", " Using Global Scaling Bounds: (1.3187587626244128, 15.69897000433602)\n", "āœ… BALMProjectionHead initialized with embedding_size: 1280\n", "āœ… BALMForLoRAFinetuning model initialized.\n", " Optimizing 2683392 trainable parameters.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 1/30 | Train Loss: 0.0359 | Test RMSE: 2.0448 | Test Pearson: 0.5113\n", " -> New best model saved with Val RMSE: 2.0448\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 2/30 | Train Loss: 0.0248 | Test RMSE: 2.0102 | Test Pearson: 0.5334\n", " -> New best model saved with Val RMSE: 2.0102\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 3/30 | Train Loss: 0.0207 | Test RMSE: 2.0488 | Test Pearson: 0.5215\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 4/30 | Train Loss: 0.0178 | Test RMSE: 2.0156 | Test Pearson: 0.5136\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 5/30 | Train Loss: 0.0161 | Test RMSE: 2.0123 | Test Pearson: 0.5187\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 6/30 | Train Loss: 0.0152 | Test RMSE: 2.0143 | Test Pearson: 0.5180\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 7/30 | Train Loss: 0.0142 | Test RMSE: 2.0667 | Test Pearson: 0.5271\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 8/30 | Train Loss: 0.0133 | Test RMSE: 2.0378 | Test Pearson: 0.5296\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 9/30 | Train Loss: 0.0129 | Test RMSE: 2.0448 | Test Pearson: 0.5257\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 10/30 | Train Loss: 0.0121 | Test RMSE: 2.0341 | Test Pearson: 0.5296\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 11/30 | Train Loss: 0.0116 | Test RMSE: 2.0265 | Test Pearson: 0.5359\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 12/30 | Train Loss: 0.0111 | Test RMSE: 1.9949 | Test Pearson: 0.5378\n", " -> New best model saved with Val RMSE: 1.9949\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 13/30 | Train Loss: 0.0109 | Test RMSE: 1.9885 | Test Pearson: 0.5425\n", " -> New best model saved with Val RMSE: 1.9885\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 14/30 | Train Loss: 0.0107 | Test RMSE: 1.9841 | Test Pearson: 0.5458\n", " -> New best model saved with Val RMSE: 1.9841\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 15/30 | Train Loss: 0.0103 | Test RMSE: 2.0208 | Test Pearson: 0.5283\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 16/30 | Train Loss: 0.0102 | Test RMSE: 2.0032 | Test Pearson: 0.5443\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 17/30 | Train Loss: 0.0096 | Test RMSE: 1.9748 | Test Pearson: 0.5406\n", " -> New best model saved with Val RMSE: 1.9748\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 18/30 | Train Loss: 0.0095 | Test RMSE: 2.0147 | Test Pearson: 0.5257\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 19/30 | Train Loss: 0.0093 | Test RMSE: 1.9866 | Test Pearson: 0.5439\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 20/30 | Train Loss: 0.0093 | Test RMSE: 1.9855 | Test Pearson: 0.5359\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 21/30 | Train Loss: 0.0089 | Test RMSE: 1.9813 | Test Pearson: 0.5390\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 22/30 | Train Loss: 0.0086 | Test RMSE: 1.9871 | Test Pearson: 0.5388\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 23/30 | Train Loss: 0.0085 | Test RMSE: 1.9775 | Test Pearson: 0.5431\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 24/30 | Train Loss: 0.0081 | Test RMSE: 1.9819 | Test Pearson: 0.5413\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 25/30 | Train Loss: 0.0080 | Test RMSE: 2.0048 | Test Pearson: 0.5376\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 26/30 | Train Loss: 0.0077 | Test RMSE: 1.9702 | Test Pearson: 0.5508\n", " -> New best model saved with Val RMSE: 1.9702\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 27/30 | Train Loss: 0.0074 | Test RMSE: 2.0032 | Test Pearson: 0.5367\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 28/30 | Train Loss: 0.0074 | Test RMSE: 1.9960 | Test Pearson: 0.5411\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 29/30 | Train Loss: 0.0073 | Test RMSE: 1.9869 | Test Pearson: 0.5361\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 30/30 | Train Loss: 0.0077 | Test RMSE: 2.0365 | Test Pearson: 0.5335\n", "Fold 3 training complete.\n", "Loading best model (based on lowest train loss) for final TEST evaluation...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\hs494\\AppData\\Local\\Temp\\ipykernel_46900\\1832587606.py:744: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " model.load_state_dict(torch.load(fold_model_path))\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Fold 3 Final Test Metrics:\n", " - RMSE: 1.9702\n", " - PEARSON: 0.5508\n", " - SPEARMAN: 0.5622\n", " - CI: 0.6971\n", " šŸ’¾ Predictions for fold 3 saved to 'cv_results_lora_seqsim\\fold_3_predictions.csv'\n", "\n", "===== STRATEGY: SequenceSimilarity | FOLD 4/5 =====\n", "šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold 4...\n", "šŸ”§ Initializing ProteinEmbeddingExtractor on cuda\n", "šŸ“„ Loading ESM-2 model: facebook/esm2_t33_650M_UR50D with dtype: torch.float16\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n", "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "āœØ Applying LoRA with rank=8, alpha=16, dropout=0.1\n", " āœ… LoRA model loaded and configured. Trainable parameters:\n", "trainable params: 2,027,520 || all params: 654,381,461 || trainable%: 0.3098\n", " āœ… Model loaded. Embedding size from config: 1280\n", " Total samples allocated: 12019\n", " Train: 7331 samples\n", " Test: 4688 samples\n", " Using Global Scaling Bounds: (1.3187587626244128, 15.69897000433602)\n", "āœ… BALMProjectionHead initialized with embedding_size: 1280\n", "āœ… BALMForLoRAFinetuning model initialized.\n", " Optimizing 2683392 trainable parameters.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 1/30 | Train Loss: 0.0470 | Test RMSE: 1.5026 | Test Pearson: 0.7055\n", " -> New best model saved with Val RMSE: 1.5026\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 2/30 | Train Loss: 0.0331 | Test RMSE: 1.4531 | Test Pearson: 0.7354\n", " -> New best model saved with Val RMSE: 1.4531\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 3/30 | Train Loss: 0.0265 | Test RMSE: 1.5033 | Test Pearson: 0.7060\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 4/30 | Train Loss: 0.0230 | Test RMSE: 1.5103 | Test Pearson: 0.7024\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 5/30 | Train Loss: 0.0200 | Test RMSE: 1.5041 | Test Pearson: 0.7084\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 6/30 | Train Loss: 0.0182 | Test RMSE: 1.5501 | Test Pearson: 0.6930\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 7/30 | Train Loss: 0.0165 | Test RMSE: 1.5200 | Test Pearson: 0.7031\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 8/30 | Train Loss: 0.0156 | Test RMSE: 1.5369 | Test Pearson: 0.7096\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 9/30 | Train Loss: 0.0145 | Test RMSE: 1.5344 | Test Pearson: 0.6938\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 10/30 | Train Loss: 0.0139 | Test RMSE: 1.5293 | Test Pearson: 0.6993\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 11/30 | Train Loss: 0.0131 | Test RMSE: 1.5106 | Test Pearson: 0.7039\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 12/30 | Train Loss: 0.0131 | Test RMSE: 1.4854 | Test Pearson: 0.7157\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 13/30 | Train Loss: 0.0124 | Test RMSE: 1.5316 | Test Pearson: 0.7007\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 14/30 | Train Loss: 0.0120 | Test RMSE: 1.5182 | Test Pearson: 0.7025\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 15/30 | Train Loss: 0.0115 | Test RMSE: 1.5174 | Test Pearson: 0.7017\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 16/30 | Train Loss: 0.0111 | Test RMSE: 1.5635 | Test Pearson: 0.6909\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 17/30 | Train Loss: 0.0110 | Test RMSE: 1.5489 | Test Pearson: 0.6935\n", "ā³ Early stopping triggered at epoch 17. Val RMSE has not improved for 15 epochs.\n", "Fold 4 training complete.\n", "Loading best model (based on lowest train loss) for final TEST evaluation...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\hs494\\AppData\\Local\\Temp\\ipykernel_46900\\1832587606.py:744: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " model.load_state_dict(torch.load(fold_model_path))\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Fold 4 Final Test Metrics:\n", " - RMSE: 1.4531\n", " - PEARSON: 0.7354\n", " - SPEARMAN: 0.7166\n", " - CI: 0.7585\n", " šŸ’¾ Predictions for fold 4 saved to 'cv_results_lora_seqsim\\fold_4_predictions.csv'\n", "\n", "===== STRATEGY: SequenceSimilarity | FOLD 5/5 =====\n", "šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold 5...\n", "šŸ”§ Initializing ProteinEmbeddingExtractor on cuda\n", "šŸ“„ Loading ESM-2 model: facebook/esm2_t33_650M_UR50D with dtype: torch.float16\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Some weights of EsmModel were not initialized from the model checkpoint at facebook/esm2_t33_650M_UR50D and are newly initialized: ['esm.pooler.dense.bias', 'esm.pooler.dense.weight']\n", "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "āœØ Applying LoRA with rank=8, alpha=16, dropout=0.1\n", " āœ… LoRA model loaded and configured. Trainable parameters:\n", "trainable params: 2,027,520 || all params: 654,381,461 || trainable%: 0.3098\n", " āœ… Model loaded. Embedding size from config: 1280\n", " Total samples allocated: 12019\n", " Train: 9397 samples\n", " Test: 2622 samples\n", " Using Global Scaling Bounds: (1.3187587626244128, 15.69897000433602)\n", "āœ… BALMProjectionHead initialized with embedding_size: 1280\n", "āœ… BALMForLoRAFinetuning model initialized.\n", " Optimizing 2683392 trainable parameters.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 1/30 | Train Loss: 0.0434 | Test RMSE: 1.5745 | Test Pearson: 0.6054\n", " -> New best model saved with Val RMSE: 1.5745\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 2/30 | Train Loss: 0.0316 | Test RMSE: 1.5507 | Test Pearson: 0.6019\n", " -> New best model saved with Val RMSE: 1.5507\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 3/30 | Train Loss: 0.0263 | Test RMSE: 1.5797 | Test Pearson: 0.5909\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 4/30 | Train Loss: 0.0229 | Test RMSE: 1.6157 | Test Pearson: 0.5819\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 5/30 | Train Loss: 0.0205 | Test RMSE: 1.6150 | Test Pearson: 0.5832\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 6/30 | Train Loss: 0.0185 | Test RMSE: 1.6439 | Test Pearson: 0.5672\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 7/30 | Train Loss: 0.0169 | Test RMSE: 1.6455 | Test Pearson: 0.5716\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 8/30 | Train Loss: 0.0159 | Test RMSE: 1.6638 | Test Pearson: 0.5578\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 9/30 | Train Loss: 0.0149 | Test RMSE: 1.6216 | Test Pearson: 0.5831\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 10/30 | Train Loss: 0.0144 | Test RMSE: 1.5911 | Test Pearson: 0.5882\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 11/30 | Train Loss: 0.0137 | Test RMSE: 1.6196 | Test Pearson: 0.5650\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 12/30 | Train Loss: 0.0129 | Test RMSE: 1.5879 | Test Pearson: 0.6014\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 13/30 | Train Loss: 0.0123 | Test RMSE: 1.6234 | Test Pearson: 0.5847\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 14/30 | Train Loss: 0.0118 | Test RMSE: 1.6296 | Test Pearson: 0.5762\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 15/30 | Train Loss: 0.0114 | Test RMSE: 1.5724 | Test Pearson: 0.5994\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 16/30 | Train Loss: 0.0112 | Test RMSE: 1.6025 | Test Pearson: 0.5816\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Epoch 17/30 | Train Loss: 0.0107 | Test RMSE: 1.6086 | Test Pearson: 0.5830\n", "ā³ Early stopping triggered at epoch 17. Val RMSE has not improved for 15 epochs.\n", "Fold 5 training complete.\n", "Loading best model (based on lowest train loss) for final TEST evaluation...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\hs494\\AppData\\Local\\Temp\\ipykernel_46900\\1832587606.py:744: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " model.load_state_dict(torch.load(fold_model_path))\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Fold 5 Final Test Metrics:\n", " - RMSE: 1.5507\n", " - PEARSON: 0.6019\n", " - SPEARMAN: 0.5709\n", " - CI: 0.7011\n", " šŸ’¾ Predictions for fold 5 saved to 'cv_results_lora_seqsim\\fold_5_predictions.csv'\n", "\n", "šŸ“ˆ Generating combined regression plot for SequenceSimilarity\n", "Plot saved to cv_results_lora_seqsim\\overall_regression_plot.png\n", "\n", "āœ… Final Results Summary for SequenceSimilarity with LoRA Fine-tuning:\n", " Mean Std Dev\n", "rmse 1.672272 0.195301\n", "pearson 0.612109 0.086380\n", "spearman 0.601683 0.086626\n", "ci 0.713320 0.034457\n", "Overall CV summary metrics saved to cv_results_lora_seqsim\\cv_summary_metrics.csv\n", "Fold model checkpoints and prediction files are saved in the 'cv_results_lora_seqsim' directory.\n", "\n", "================================================================================\n", "... 8. FINAL NOTES AND RECOMMENDATIONS ...\n", "\n", "šŸŽ‰ CONGRATULATIONS! You've successfully implemented BALM with on-the-fly ESM-2 embeddings\n", "and **Parameter-Efficient Fine-Tuning (PEFT) using LoRA**!\n", "\n", "šŸ“Š KEY ADVANTAGES OF THIS APPROACH:\n", "ā€¢ **Efficiency:** Fine-tunes only a tiny fraction of the ESM-2 model parameters (<1%).\n", "ā€¢ **Speed & Memory:** Significantly faster and less memory-intensive than full fine-tuning.\n", "ā€¢ **Performance:** Adapts the powerful ESM-2 model to your specific data, often leading to better performance.\n", "ā€¢ **No Static Cache:** By computing embeddings dynamically, the model fully leverages the fine-tuned LoRA adapters during training.\n", "\n", "šŸš€ PERFORMANCE OPTIMIZATION TIPS:\n", "1. **LoRA Hyperparameters:**\n", " * `lora_rank (r)`: Controls the capacity of the adapters. Higher `r` means more parameters. Typical values are 8-64.\n", " * `lora_alpha`: Acts as a scaling factor. A common heuristic is to set `lora_alpha` to be twice the `lora_rank`. [7]\n", " * `target_modules`: Experiment with adding other modules like `dense` to the target list in the `LoraConfig`. [1]\n", "2. **Learning Rate:** LoRA can sometimes handle slightly higher learning rates than full fine-tuning. Experiment with values from 5e-5 to 5e-4.\n", "3. **Batch Size & Gradient Accumulation:** If your GPU memory is a constraint, you can simulate a larger batch size by using gradient accumulation.\n", "\n", "šŸ’¾ SAVING AND LOADING THE FINE-TUNED MODEL:\n", "ā€¢ After training, you can save the trained LoRA adapters separately from the base model. This is very efficient for storage. [3]\n", "ā€¢ To save: `model.esm_model.save_pretrained(\"path/to/my_lora_adapters\")` [3]\n", "ā€¢ To load: First, load the base ESM-2 model. Then, use `PeftModel.from_pretrained(base_model, \"path/to/my_lora_adapters\")`.\n", "\n", "šŸŽÆ NEXT STEPS:\n", "1. **Hyperparameter Search:** Systematically tune LoRA parameters (`r`, `lora_alpha`) and the learning rate for optimal performance.\n", "2. **Experiment with Larger ESM-2 Models:** LoRA makes fine-tuning even the largest ESM-2 models (e.g., 3B parameters) feasible on a single high-end GPU.\n", "3. **Advanced PEFT Techniques:** Explore other PEFT methods like IA3 or full parameter-efficient fine-tuning for comparison.\n", "\n", "\n", "āœ… LoRA PEFT implementation complete! Happy parameter-efficient protein modeling! šŸ§¬šŸ¤–\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\hs494\\AppData\\Local\\Temp\\ipykernel_46900\\1832587606.py:423: UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown\n", " plt.show()\n" ] } ], "source": [ "\n", "import os\n", "os.environ.pop('MPLBACKEND', None)\n", "\n", "os.environ['MPLBACKEND'] = 'agg'\n", "\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "\n", "\n", "# ==============================================================================\n", "# 0. PEFT IMPORTS AND INITIAL SETUP (NEW)\n", "# ==============================================================================\n", "# Ensure you have the peft library installed: pip install peft\n", "from peft import LoraConfig, get_peft_model, TaskType\n", "\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 0. PEFT LIBRARY INITIALIZED ...\")\n", "print(\" Ready for Parameter-Efficient Fine-Tuning with LoRA!\")\n", "# ==============================================================================\n", "# 1. IMPORTS AND INITIAL SETUP\n", "# ==============================================================================\n", "import os\n", "import sys\n", "import time\n", "import json\n", "import random\n", "import pickle\n", "from pathlib import Path\n", "from datetime import datetime\n", "\n", "# Standard data science libraries\n", "import numpy as np\n", "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "# In the IMPORTS AND INITIAL SETUP section\n", "from tqdm import tqdm\n", "import warnings\n", "\n", "# Machine learning libraries\n", "from sklearn.model_selection import KFold, GroupKFold\n", "from sklearn.metrics import mean_squared_error\n", "from scipy.stats import pearsonr, spearmanr\n", "from sklearn.cluster import AgglomerativeClustering\n", "\n", "\n", "# PyTorch and Transformers\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torch.utils.data import Dataset, DataLoader\n", "from torch.optim import AdamW\n", "from transformers import AutoModel, AutoTokenizer\n", "\n", "# Advanced CV imports\n", "from Bio.Align import PairwiseAligner\n", "from scipy.cluster.hierarchy import linkage, fcluster\n", "from multiprocessing import Pool, cpu_count\n", "from typing import List, Tuple\n", "\n", "# Set style for plots\n", "\n", "\n", "# --- MODIFIED: Title reflects the use of ESM-2 ---\n", "print(\"šŸ§¬ BALM Adaptation for ESM-2 with LoRA Fine-Tuning\")\n", "print(\"āš” Fast Training with 5-Fold Cross-Validation\")\n", "print(\"=\"*80)\n", "\n", "def setup_reproducibility(seed=42):\n", " \"\"\"Setup complete reproducibility.\"\"\"\n", " print(f\"šŸ”’ Setting up reproducibility with seed {seed}\")\n", " random.seed(seed)\n", " np.random.seed(seed)\n", " torch.manual_seed(seed)\n", " if torch.cuda.is_available():\n", " torch.cuda.manual_seed(seed)\n", " torch.cuda.manual_seed_all(seed)\n", " torch.backends.cudnn.deterministic = True\n", " torch.backends.cudnn.benchmark = False\n", " os.environ['PYTHONHASHSEED'] = str(seed)\n", " print(f\"āœ… Reproducibility setup complete\")\n", "\n", "setup_reproducibility(42)\n", "\n", "# ==============================================================================\n", "# 2. ESM-2 EMBEDDING EXTRACTOR (Modified for LoRA)\n", "# ==============================================================================\n", "\n", "class ProteinEmbeddingExtractor:\n", " \"\"\"\n", " A class to prepare the ESM-2 model, optionally with LoRA adapters.\n", " This class handles model loading and tokenization.\n", " \"\"\"\n", " def __init__(\n", " self,\n", " model_name=\"facebook/esm2_t33_650M_UR50D\",\n", " device=\"auto\",\n", " lora_rank=8,\n", " lora_alpha=16,\n", " lora_dropout=0.1,\n", " use_lora=True\n", " ):\n", " self.model_name = model_name\n", " self.use_lora = use_lora\n", " \n", " if device == \"auto\":\n", " self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " else:\n", " self.device = torch.device(device)\n", " \n", " print(f\"šŸ”§ Initializing ProteinEmbeddingExtractor on {self.device}\")\n", " \n", " self.dtype = torch.float16 if self.device.type == 'cuda' else torch.float32\n", " print(f\"šŸ“„ Loading ESM-2 model: {model_name} with dtype: {self.dtype}\")\n", " \n", " self.model = AutoModel.from_pretrained(\n", " model_name,\n", " torch_dtype=self.dtype\n", " )\n", " self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n", " \n", " if self.use_lora:\n", " print(f\"āœØ Applying LoRA with rank={lora_rank}, alpha={lora_alpha}, dropout={lora_dropout}\")\n", " lora_config = LoraConfig(\n", " r=lora_rank,\n", " lora_alpha=lora_alpha,\n", " lora_dropout=lora_dropout,\n", " bias=\"none\",\n", " task_type=TaskType.FEATURE_EXTRACTION,\n", " target_modules=[\"key\", \"query\", \"value\"]\n", " )\n", " self.model = get_peft_model(self.model, lora_config)\n", " print(\" āœ… LoRA model loaded and configured. Trainable parameters:\")\n", " self.model.print_trainable_parameters()\n", " else:\n", " print(\" āš ļø LoRA is disabled. All base model weights are frozen.\")\n", " for param in self.model.parameters():\n", " param.requires_grad = False\n", "\n", " self.model.to(self.device)\n", " self.embedding_size = self.model.config.hidden_size\n", " print(f\" āœ… Model loaded. Embedding size from config: {self.embedding_size}\")\n", "\n", " def get_model_and_tokenizer(self):\n", " \"\"\"Returns the prepared model (PeftModel or base) and tokenizer.\"\"\"\n", " return self.model, self.tokenizer\n", "\n", "# ==============================================================================\n", "# 3. LOAD YOUR DATASET\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 3. LOADING DATA FROM Data.csv ...\")\n", "\n", "protein1_col = 'Target'\n", "protein2_col = 'proteina'\n", "target_col = 'Y'\n", "pdb_col_raw = 'PDB'\n", "subgroup_col = 'Subgroup'\n", "source_dataset_col = 'Source Data Set' # NEW\n", "try:\n", " df_raw = pd.read_csv(r\"C:\\Users\\hs494\\PPIBALM\\BALM-PPI\\scripts\\notebooks\\PPB_Affinity_Sequences_Final (version 1).csv\")\n", " print(f\"Successfully loaded Data.csv with {len(df_raw)} rows.\")\n", " df = df_raw[[protein1_col, protein2_col, target_col, pdb_col_raw, subgroup_col,source_dataset_col]].copy()\n", " df.rename(columns={\n", " protein1_col: 'Target',\n", " protein2_col: 'proteina',\n", " target_col: 'Y',\n", " subgroup_col : 'Subgroup',\n", " source_dataset_col: 'SourceDataSet' # NEW\n", " }, inplace=True)\n", " \n", " print(\"\\n--- Initial Data Check ---\")\n", " print(df.head())\n", " missing_values = df.isnull().sum()\n", " print(\"\\nMissing values per column:\\n\", missing_values)\n", " \n", " if missing_values.sum() > 0:\n", " print(f\"\\nWarning: Found {missing_values.sum()} total missing values. Dropping rows with NaNs.\")\n", " df.dropna(subset=['Target', 'proteina', 'Y', 'PDB'], inplace=True)\n", " print(f\"Dataset now has {len(df)} rows after cleaning.\")\n", " \n", " df['Target'] = df['Target'].astype(str)\n", " df['proteina'] = df['proteina'].astype(str)\n", " df['Y'] = pd.to_numeric(df['Y'], errors='coerce')\n", " df['PDB'] = df['PDB'].astype(str)\n", " df['Subgroup'] = df['Subgroup'].astype(str)\n", " df['SourceDataSet'] = df['SourceDataSet'].astype(str) # NEW\n", "\n", " df.dropna(subset=['Y'], inplace=True)\n", " \n", " print(f\"\\nāœ… Data is loaded and cleaned. Final dataset size: {len(df)} samples.\")\n", " df.info()\n", "\n", "except FileNotFoundError:\n", " print(\"āŒ ERROR: 'Data.csv' not found.\")\n", " df = pd.DataFrame(columns=['Target', 'proteina', 'Y', 'PDB'])\n", "except KeyError as e:\n", " print(f\"āŒ ERROR: Column {e} not found in 'Data.csv'.\")\n", " df = pd.DataFrame(columns=['Target', 'proteina', 'Y', 'PDB'])\n", "\n", "# ==============================================================================\n", "# 4. EXPERIMENT CONFIGURATION\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 4. CONFIGURING THE EXPERIMENT ...\")\n", "\n", "CV_STRATEGY = 'SequenceSimilarity' \n", "N_SPLITS = 5\n", "EPOCHS = 30\n", "PATIENCE = 15\n", "BATCH_SIZE = 1\n", "LEARNING_RATE = 1e-4\n", "PROJECTED_SIZE = 256\n", "DROPOUT = 0.1\n", "\n", "LORA_RANK = 8\n", "LORA_ALPHA = 16\n", "LORA_DROPOUT = 0.1\n", "USE_LORA_FOR_TRAINING = True\n", "\n", "# Sequence Similarity CV Configs\n", "SIMILARITY_THRESHOLD = 0.3\n", "K_MER_SIZE = 3\n", "CLUSTERING_METHOD = 'agglomerative'\n", "\n", "RESULTS_DIR = \"cv_results_lora_seqsim\" # Use a unique name\n", "os.makedirs(RESULTS_DIR, exist_ok=True)\n", "print(f\"Results will be saved in: {RESULTS_DIR}\")\n", "\n", "print(f\"āœ… Configuration set. Using CV Strategy: {CV_STRATEGY}\")\n", "print(f\" Epochs: {EPOCHS}, Patience: {PATIENCE}\")\n", "if USE_LORA_FOR_TRAINING:\n", " print(f\" LoRA Enabled: Rank={LORA_RANK}, Alpha={LORA_ALPHA}, Dropout={LORA_DROPOUT}\")\n", "else:\n", " print(f\" LoRA Disabled: Only projection head will be trained.\")\n", "\n", "# ==============================================================================\n", "# 5. ADAPTED MODEL AND DATASET\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 5. DEFINING ADAPTED MODEL AND DATASET ...\")\n", "\n", "class BALMProjectionHead(nn.Module):\n", " def __init__(self, embedding_size, projected_size, projected_dropout):\n", " super().__init__()\n", " self.protein_projection = nn.Linear(embedding_size, projected_size)\n", " self.proteina_projection = nn.Linear(embedding_size, projected_size)\n", " self.dropout = nn.Dropout(projected_dropout)\n", " self.loss_fn = nn.MSELoss()\n", " print(f\"āœ… BALMProjectionHead initialized with embedding_size: {embedding_size}\")\n", "\n", " def forward(self, protein_embedding, proteina_embedding, labels=None):\n", " protein_embedding = self.dropout(protein_embedding)\n", " proteina_embedding = self.dropout(proteina_embedding)\n", " protein_projected = self.protein_projection(protein_embedding)\n", " proteina_projected = self.proteina_projection(proteina_embedding)\n", " protein_projected = F.normalize(protein_projected, p=2, dim=1)\n", " proteina_projected = F.normalize(proteina_projected, p=2, dim=1)\n", " cosine_similarity = F.cosine_similarity(protein_projected, proteina_projected)\n", " cosine_similarity = torch.clamp(cosine_similarity, -0.9999, 0.9999)\n", " output = {\"cosine_similarity\": cosine_similarity}\n", " if labels is not None:\n", " loss = self.loss_fn(cosine_similarity, labels)\n", " output[\"loss\"] = loss\n", " return output\n", "\n", "class BALMForLoRAFinetuning(nn.Module):\n", " def __init__(self, esm_model, esm_tokenizer, projected_size, projected_dropout, pkd_bounds):\n", " super().__init__()\n", " self.esm_model = esm_model\n", " self.esm_tokenizer = esm_tokenizer\n", " self.embedding_size = self.esm_model.config.hidden_size\n", " self.projection_head = BALMProjectionHead(self.embedding_size, projected_size, projected_dropout)\n", " self.pkd_lower, self.pkd_upper = pkd_bounds\n", " self.pkd_range = self.pkd_upper - self.pkd_lower\n", " print(f\"āœ… BALMForLoRAFinetuning model initialized.\")\n", "\n", " def _get_esm_embeddings(self, sequences: List[str]):\n", " cls_token = self.esm_tokenizer.cls_token\n", " processed_seqs = [s.replace('|', f\"{cls_token}{cls_token}\") for s in sequences]\n", " \n", " inputs = self.esm_tokenizer(processed_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024)\n", " inputs = {k: v.to(self.esm_model.device) for k, v in inputs.items()}\n", " outputs = self.esm_model(**inputs)\n", " \n", " # Pooling logic (remains Mean Pooling, which now covers the internal CLS tokens)\n", " hidden_states = outputs.last_hidden_state\n", " mask = inputs['attention_mask'].unsqueeze(-1).expand(hidden_states.size()).float()\n", " return torch.sum(hidden_states * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)\n", "\n", " def forward(self, batch_input):\n", " protein_emb = self._get_esm_embeddings(batch_input[\"protein_sequence\"])\n", " proteina_emb = self._get_esm_embeddings(batch_input[\"proteina_sequence\"])\n", " labels = batch_input.get(\"labels\")\n", " proj_output = self.projection_head(protein_emb, proteina_emb, labels)\n", " return {\n", " \"cosine_similarity\": proj_output[\"cosine_similarity\"],\n", " \"original_pkds\": batch_input[\"original_pkds\"],\n", " \"pdb_groups\": batch_input[\"pdb_groups\"],\n", " \"subgroups\": batch_input[\"subgroups\"],\n", " \"source_dataset\": batch_input[\"source_dataset\"], # NEW\n", " \"loss\": proj_output.get(\"loss\")\n", " }\n", " if \"loss\" in proj_output:\n", "\n", " output[\"loss\"] = proj_output[\"loss\"]\n", " return output\n", "\n", "class ProteinPairSequenceDataset(Dataset):\n", " def __init__(self, dataframe, pkd_bounds):\n", " self.data = dataframe.reset_index(drop=True)\n", " self.pkd_lower, self.pkd_upper = pkd_bounds\n", " self.pkd_range = self.pkd_upper - self.pkd_lower\n", "\n", " def __len__(self):\n", " return len(self.data)\n", "\n", " def __getitem__(self, idx):\n", " row = self.data.iloc[idx]\n", " original_pkd = float(row['Y'])\n", " scaled_label = ((original_pkd - self.pkd_lower) / self.pkd_range) * 2 - 1\n", " return {\n", " \"protein_sequence\": row[\"Target\"],\n", " \"proteina_sequence\": row[\"proteina\"],\n", " \"labels\": torch.tensor(scaled_label, dtype=torch.float32),\n", " \"original_pkds\": torch.tensor(original_pkd, dtype=torch.float32),\n", " \"pdb_groups\": row[\"PDB\"],\n", " \"subgroups\": row[\"Subgroup\"],\n", " \"source_dataset\": row[\"SourceDataSet\"] # NEW\n", " }\n", "\n", "def collate_fn_sequences(batch):\n", " return {\n", " \"protein_sequence\": [item['protein_sequence'] for item in batch],\n", " \"proteina_sequence\": [item['proteina_sequence'] for item in batch],\n", " \"labels\": torch.stack([item['labels'] for item in batch]),\n", " \"original_pkds\": torch.stack([item['original_pkds'] for item in batch]),\n", " \"pdb_groups\": [item['pdb_groups'] for item in batch], # NEW\n", " \"subgroups\": [item['subgroups'] for item in batch], # NEW\n", " \"source_dataset\": [item['source_dataset'] for item in batch] # NEW\n", " }\n", "\n", "# ==============================================================================\n", "# 6. METRICS, EVALUATION & SEQUENCE SIMILARITY HELPERS (NEW)\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 6. DEFINING METRICS, EVALUATION & SEQUENCE HELPERS ...\")\n", "\n", "def concordance_index(y_true, y_pred):\n", " try:\n", " y_true, y_pred = np.array(y_true), np.array(y_pred)\n", " n = len(y_true)\n", " if n < 2: return 1.0\n", " concordant, discordant, tied_pred = 0, 0, 0\n", " for i in range(n):\n", " for j in range(i + 1, n):\n", " if y_true[i] == y_true[j]: continue\n", " if y_pred[i] == y_pred[j]:\n", " tied_pred += 1\n", " continue\n", " if (y_pred[i] > y_pred[j] and y_true[i] > y_true[j]) or \\\n", " (y_pred[i] < y_pred[j] and y_true[i] < y_true[j]):\n", " concordant += 1\n", " else:\n", " discordant += 1\n", " total_pairs = concordant + discordant + tied_pred\n", " return 1.0 if total_pairs == 0 else (concordant + 0.5 * tied_pred) / total_pairs\n", " except Exception:\n", " return np.nan\n", "\n", "def evaluate_model(model, data_loader, device, pkd_bounds):\n", " model.eval()\n", " all_labels, all_preds_pkd, all_pdb_groups, all_subgroups, all_sources = [], [], [], [], [] \n", " pkd_lower, pkd_range = pkd_bounds[0], pkd_bounds[1] - pkd_bounds[0]\n", " \n", " with torch.no_grad():\n", " for batch in data_loader:\n", " batch_gpu = {k: v for k, v in batch.items()}\n", " for key in [\"labels\", \"original_pkds\"]:\n", " if key in batch_gpu: batch_gpu[key] = batch_gpu[key].to(device)\n", " outputs = model(batch_gpu)\n", " cosine_sim = outputs['cosine_similarity'].cpu()\n", " preds_pkd = ((cosine_sim + 1) / 2) * pkd_range + pkd_lower\n", " all_labels.extend(batch['original_pkds'].numpy())\n", " all_preds_pkd.extend(preds_pkd.numpy())\n", " all_pdb_groups.extend(batch['pdb_groups']) \n", " all_subgroups.extend(batch['subgroups']) \n", " all_sources.extend(batch['source_dataset']) # NEW\n", "\n", " labels, preds = np.array(all_labels), np.array(all_preds_pkd)\n", " pdb_groups_arr = np.array(all_pdb_groups)\n", " subgroups_arr = np.array(all_subgroups)\n", " sources_arr = np.array(all_sources) # NEW\n", " \n", " mask = ~np.isnan(labels) & ~np.isnan(preds)\n", " labels, preds = labels[mask], preds[mask]\n", " \n", " if len(labels) < 2:\n", " return {\"rmse\": np.nan, \"pearson\": np.nan, \"spearman\": np.nan, \"ci\": np.nan}, labels, preds, pdb_groups_arr, subgroups_arr, sources_arr\n", " return {\n", " 'rmse': np.sqrt(mean_squared_error(labels, preds)),\n", " 'pearson': pearsonr(labels, preds)[0],\n", " 'spearman': spearmanr(labels, preds)[0],\n", " 'ci': concordance_index(labels, preds)\n", " }, labels, preds, pdb_groups_arr, subgroups_arr, sources_arr\n", "\n", "def plot_regression(y_true, y_pred, metrics, title, filename=None):\n", " plt.figure(figsize=(8, 8))\n", " sns.scatterplot(x=y_true, y=y_pred, alpha=0.6, edgecolor='k', s=50)\n", " min_val, max_val = min(np.min(y_true), np.min(y_pred)), max(np.max(y_true), np.max(y_pred))\n", " plt.plot([min_val, max_val], [min_val, max_val], 'r--', lw=2, label='Identity Line (y=x)')\n", " metrics_text = f\"RMSE: {metrics['rmse']:.4f}\\nPearson: {metrics['pearson']:.4f}\\nSpearman: {metrics['spearman']:.4f}\\nCI: {metrics['ci']:.4f}\"\n", " plt.text(0.05, 0.95, metrics_text, transform=plt.gca().transAxes, fontsize=12, verticalalignment='top', bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.5))\n", " plt.title(title, fontsize=16)\n", " plt.xlabel(\"Actual pKd\", fontsize=14)\n", " plt.ylabel(\"Predicted pKd\", fontsize=14)\n", " plt.legend()\n", " plt.grid(True)\n", " plt.tight_layout()\n", " if filename: # NEW\n", " plt.savefig(filename) # NEW\n", " print(f\"Plot saved to {filename}\") # NEW\n", " plt.show()\n", "\n", "def compute_kmer_similarity_matrix(sequences, k=3):\n", " \"\"\"\n", " Compute pairwise sequence similarity matrix using k-mer based Jaccard similarity.\n", " This is a biologically relevant method for sequence similarity computation.\n", " \"\"\"\n", " print(f\"Computing {k}-mer Jaccard similarity matrix for {len(sequences)} sequences...\")\n", " \n", " def get_kmers(sequence, k_size):\n", " \"\"\"Extract k-mers from a sequence.\"\"\"\n", " if len(sequence) < k_size:\n", " return {sequence}\n", " return {sequence[i:i+k_size] for i in range(len(sequence) - k_size + 1)}\n", " \n", " def jaccard_similarity(seq1, seq2):\n", " \"\"\"Compute Jaccard similarity between two sequences based on k-mers.\"\"\"\n", " kmers1 = get_kmers(seq1, k)\n", " kmers2 = get_kmers(seq2, k)\n", " if not kmers1 or not kmers2:\n", " return 0.0\n", " intersection = len(kmers1 & kmers2)\n", " union = len(kmers1 | kmers2)\n", " return intersection / union if union > 0 else 0.0\n", " \n", " n_seqs = len(sequences)\n", " similarity_matrix = np.zeros((n_seqs, n_seqs))\n", " \n", " # Compute similarity matrix with progress bar\n", " for i in tqdm(range(n_seqs), desc=\"Computing similarities\"):\n", " for j in range(i, n_seqs):\n", " if i == j:\n", " similarity_matrix[i, j] = 1.0\n", " else:\n", " sim = jaccard_similarity(sequences[i], sequences[j])\n", " similarity_matrix[i, j] = sim\n", " similarity_matrix[j, i] = sim\n", " \n", " avg_similarity = np.mean(similarity_matrix[np.triu_indices_from(similarity_matrix, k=1)])\n", " print(f\"Average pairwise similarity: {avg_similarity:.3f}\")\n", " \n", " return similarity_matrix\n", "\n", "def cluster_sequences_by_similarity(similarity_matrix, threshold=0.3, method='agglomerative'):\n", " \"\"\"\n", " Cluster sequences based on similarity matrix using specified clustering method.\n", " Uses distance threshold derived from similarity threshold.\n", " MODIFIED: Uses scipy.cluster.hierarchy for deterministic clustering.\n", " \"\"\"\n", " print(f\"Clustering sequences using {method} method with similarity threshold {threshold}\")\n", " \n", " # Convert similarity to distance matrix\n", " distance_matrix = 1.0 - similarity_matrix\n", " \n", " try:\n", " if method == 'agglomerative':\n", " from scipy.spatial.distance import squareform\n", " from scipy.cluster.hierarchy import linkage, fcluster\n", " \n", " # Use distance threshold - sequences with similarity > threshold will be in same cluster\n", " distance_threshold = 1.0 - threshold\n", " \n", " print(f\"Using distance threshold: {distance_threshold}\")\n", " print(f\"Distance matrix stats - Min: {distance_matrix.min():.3f}, Max: {distance_matrix.max():.3f}, Mean: {distance_matrix.mean():.3f}\")\n", " \n", " # Convert to condensed distance matrix (required by scipy.linkage)\n", " # squareform converts square distance matrix to condensed vector\n", " condensed_dist = squareform(distance_matrix, checks=False)\n", " \n", " # Perform hierarchical clustering (DETERMINISTIC)\n", " # method='average' corresponds to UPGMA (average linkage)\n", " Z = linkage(condensed_dist, method='average')\n", " \n", " # Cut dendrogram at distance threshold to form flat clusters\n", " # criterion='distance' cuts at specified height\n", " clusters = fcluster(Z, t=distance_threshold, criterion='distance')\n", " \n", " # Convert to 0-indexed (fcluster returns 1-indexed)\n", " clusters = clusters - 1\n", " \n", " print(f\"Hierarchical clustering completed successfully (deterministic)\")\n", " \n", " else:\n", " raise ValueError(f\"Unsupported clustering method: {method}\")\n", " \n", " if clusters is None:\n", " raise ValueError(\"Clustering failed - returned None\")\n", " \n", " n_clusters = len(set(clusters))\n", " print(f\"Created {n_clusters} clusters\")\n", " \n", " cluster_sizes = [np.sum(clusters == i) for i in set(clusters)]\n", " print(f\"Cluster size distribution - Min: {min(cluster_sizes)}, Max: {max(cluster_sizes)}, Mean: {np.mean(cluster_sizes):.1f}\")\n", " \n", " return clusters\n", " \n", " except Exception as e:\n", " print(f\"āŒ Clustering failed with error: {e}\")\n", " print(\"šŸ”„ Falling back to individual clusters (each sequence in its own cluster)\")\n", " \n", " n_sequences = similarity_matrix.shape[0]\n", " fallback_clusters = np.arange(n_sequences)\n", " \n", " print(f\"Created {n_sequences} individual clusters (fallback)\")\n", " return fallback_clusters\n", "\n", "# ==============================================================================\n", "# 7. CROSS-VALIDATION TRAINING (MODIFIED FOR 80/20 SPLIT + TRAIN LOSS PATIENCE)\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(f\"... 7. RUNNING '{CV_STRATEGY}' STRATEGY WITH LoRA (80/20 Split with Training Loss Patience) ...\")\n", "\n", "def run_sequence_similarity_cv_with_lora():\n", " if df.empty:\n", " print(\"āš ļø Skipping training due to empty dataframe.\")\n", " return\n", "\n", " device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", " PKD_LOWER_BOUND, PKD_UPPER_BOUND = df['Y'].min(), df['Y'].max()\n", " PKD_BOUNDS = (PKD_LOWER_BOUND, PKD_UPPER_BOUND)\n", " print(f\"Global pKd bounds for scaling: [{PKD_LOWER_BOUND:.2f}, {PKD_UPPER_BOUND:.2f}]\")\n", "\n", " # --- Steps 1-4: Data Preparation and Clustering ---\n", " print(\"\\nšŸ§¬ Step 1: Extracting unique protein sequences...\")\n", " sequence_columns = ['Target', 'proteina']\n", " all_sequences = set()\n", " sequence_to_records = {}\n", " for idx, row in df.iterrows():\n", " for col in sequence_columns:\n", " seq = str(row[col]).strip()\n", " if seq:\n", " all_sequences.add(seq)\n", " if seq not in sequence_to_records:\n", " sequence_to_records[seq] = []\n", " sequence_to_records[seq].append(idx)\n", " unique_sequences = sorted(list(all_sequences))\n", " print(f\"Found {len(unique_sequences)} unique protein sequences.\")\n", "\n", " print(\"\\nšŸ§¬ Step 2: Computing sequence similarity matrix...\")\n", " similarity_matrix = compute_kmer_similarity_matrix(unique_sequences, k=K_MER_SIZE)\n", "\n", " print(\"\\nšŸ§¬ Step 3: Clustering sequences...\")\n", " clusters = cluster_sequences_by_similarity(similarity_matrix, threshold=SIMILARITY_THRESHOLD, method=CLUSTERING_METHOD)\n", " \n", " cluster_to_records = {}\n", " for seq_idx, cluster_id in enumerate(clusters):\n", " seq = unique_sequences[seq_idx]\n", " if cluster_id not in cluster_to_records:\n", " cluster_to_records[cluster_id] = []\n", " cluster_to_records[cluster_id].extend(sequence_to_records[seq])\n", " \n", " for cluster_id in cluster_to_records:\n", " cluster_to_records[cluster_id] = list(set(cluster_to_records[cluster_id]))\n", " # --- End of omitted code ---\n", "\n", "\n", "\n", " # --- Step 5: Perform Cluster-based Cross-Validation ---\n", " print(f\"\\nšŸš€ Starting {N_SPLITS}-Fold **{CV_STRATEGY}** Cross-Validation with LoRA fine-tuning...\")\n", "\n", " # NEW: Display cluster count and fold split preview\n", " print(f\"\\nCreated {len(set(clusters))} clusters\")\n", " print(f\" Step 4: Creating {N_SPLITS}-fold splits...\")\n", "\n", " # Pre-compute and display split sizes\n", " cluster_ids = list(cluster_to_records.keys())\n", " kf_preview = KFold(n_splits=N_SPLITS, shuffle=True, random_state=42)\n", "\n", " for preview_fold, (train_cluster_idx, test_cluster_idx) in enumerate(kf_preview.split(cluster_ids)):\n", " preview_fold_num = preview_fold + 1\n", " \n", " train_clusters_preview = set([cluster_ids[i] for i in train_cluster_idx])\n", " test_clusters_preview = set([cluster_ids[i] for i in test_cluster_idx])\n", " \n", " sequence_to_group_preview = {}\n", " for seq_idx, cluster_id in enumerate(clusters):\n", " sequence = unique_sequences[seq_idx]\n", " if cluster_id in train_clusters_preview:\n", " sequence_to_group_preview[sequence] = 'train'\n", " elif cluster_id in test_clusters_preview:\n", " sequence_to_group_preview[sequence] = 'test'\n", " \n", " train_count, test_count = 0, 0\n", " for idx, row in df.iterrows():\n", " seq_target = str(row['Target']).strip()\n", " seq_proteina = str(row['proteina']).strip()\n", " \n", " cluster_target_id = sequence_to_group_preview.get(seq_target)\n", " cluster_proteina_id = sequence_to_group_preview.get(seq_proteina)\n", " \n", " if cluster_target_id == 'test' or cluster_proteina_id == 'test':\n", " test_count += 1\n", " else:\n", " train_count += 1\n", " \n", " print(f\" Fold {preview_fold_num}: Train={train_count}, Test={test_count}\")\n", "\n", " print(f\"\\nšŸ”„ Beginning actual training across {N_SPLITS} folds...\")\n", "\n", " \n", " kf = KFold(n_splits=N_SPLITS, shuffle=True, random_state=42)\n", " \n", " fold_results, all_folds_y_true, all_folds_y_pred = [], [], []\n", " \n", " for fold, (train_cluster_idx, test_cluster_idx) in enumerate(kf.split(cluster_ids)):\n", " fold_num = fold + 1\n", " print(f\"\\n===== STRATEGY: {CV_STRATEGY} | FOLD {fold_num}/{N_SPLITS} =====\")\n", " \n", " # --- CRITICAL FIX: RE-INITIALIZE MODEL HERE ---\n", " # This ensures the LoRA weights are reset for every fold.\n", " print(f\"šŸ”„ Initializing fresh ESM-2 Model and LoRA Adapters for Fold {fold_num}...\")\n", " extractor = ProteinEmbeddingExtractor(\n", " device=device, use_lora=USE_LORA_FOR_TRAINING, lora_rank=LORA_RANK,\n", " lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT\n", " )\n", " esm_model, esm_tokenizer = extractor.get_model_and_tokenizer()\n", " # ----------------------------------------------\n", "\n", " fold_model_path = os.path.join(RESULTS_DIR, f\"best_model_fold_{fold_num}.pth\")\n", " \n", " train_clusters = [cluster_ids[i] for i in train_cluster_idx]\n", " test_clusters = [cluster_ids[i] for i in test_cluster_idx]\n", "\n", " # Create a map from each unique sequence to its group ('train' or 'test')\n", " sequence_to_group = {}\n", " train_cluster_set = set(train_clusters)\n", " test_cluster_set = set(test_clusters)\n", "\n", " for seq_idx, cluster_id in enumerate(clusters):\n", " sequence = unique_sequences[seq_idx]\n", " if cluster_id in train_cluster_set:\n", " sequence_to_group[sequence] = 'train'\n", " elif cluster_id in test_cluster_set:\n", " sequence_to_group[sequence] = 'test'\n", "\n", " # Iterate through the main dataframe and assign rows\n", " train_indices, test_indices = [], []\n", " \n", " for idx, row in df.iterrows():\n", " seq_target = str(row['Target']).strip()\n", " seq_proteina = str(row['proteina']).strip()\n", "\n", " cluster_target_id = sequence_to_group.get(seq_target)\n", " cluster_proteina_id = sequence_to_group.get(seq_proteina)\n", "\n", " # Strict Leakage Rule: If EITHER protein is in a test cluster, it goes to test.\n", " if cluster_target_id == 'test' or cluster_proteina_id == 'test':\n", " test_indices.append(idx)\n", " else:\n", " train_indices.append(idx)\n", "\n", " # Create dataframes\n", " train_df = df.loc[train_indices]\n", " test_df = df.loc[test_indices]\n", " \n", " assert len(set(train_indices).intersection(set(test_indices))) == 0, \"Overlap found between train and test sets!\"\n", "\n", " print(f\" Total samples allocated: {len(train_df) + len(test_df)}\")\n", " print(f\" Train: {len(train_df)} samples\")\n", " print(f\" Test: {len(test_df)} samples\")\n", "\n", " # --- Use GLOBAL Scaling Bounds (for consistent comparison across folds) ---\n", " print(f\" Using Global Scaling Bounds: {PKD_BOUNDS}\")\n", "\n", " # --- DataLoaders ---\n", " train_dataset = ProteinPairSequenceDataset(train_df, PKD_BOUNDS)\n", " test_dataset = ProteinPairSequenceDataset(test_df, PKD_BOUNDS)\n", " train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn_sequences)\n", " test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE*2, shuffle=False, collate_fn=collate_fn_sequences)\n", "\n", " # --- Model, Optimizer ---\n", " model = BALMForLoRAFinetuning(\n", " esm_model=esm_model, esm_tokenizer=esm_tokenizer,\n", " projected_size=PROJECTED_SIZE, projected_dropout=DROPOUT,\n", " pkd_bounds=PKD_BOUNDS\n", " ).to(device)\n", " \n", " optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)\n", " print(f\" Optimizing {sum(p.numel() for p in model.parameters() if p.requires_grad)} trainable parameters.\")\n", "\n", " # --- Training Loop ---\n", " best_val_metric = float('inf')\n", " patience_counter = 0\n", "\n", " for epoch in range(EPOCHS):\n", " model.train()\n", " total_train_loss = 0.0\n", " progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS}\", leave=False)\n", " \n", " for batch in progress_bar:\n", " optimizer.zero_grad()\n", " batch_gpu = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}\n", " outputs = model(batch_gpu)\n", " loss = outputs['loss']\n", " loss.backward()\n", " optimizer.step()\n", " total_train_loss += loss.item()\n", " progress_bar.set_postfix({'batch_loss': f\"{loss.item():.4f}\"})\n", " \n", " avg_train_loss = total_train_loss / len(train_loader)\n", " # Evaluate on test set\n", " val_metrics, _, _, _, _, _ = evaluate_model(model, test_loader, device, PKD_BOUNDS)\n", " current_val_metric = val_metrics['rmse']\n", " print(f\"Epoch {epoch+1}/{EPOCHS} | Train Loss: {avg_train_loss:.4f} | Test RMSE: {current_val_metric:.4f} | Test Pearson: {val_metrics['pearson']:.4f}\")\n", " \n", " # Save best model based on Validation RMSE (standard practice)\n", " if current_val_metric < best_val_metric:\n", " best_val_metric = current_val_metric\n", " patience_counter = 0\n", " torch.save(model.state_dict(), fold_model_path)\n", " print(f\" -> New best model saved with Val RMSE: {best_val_metric:.4f}\")\n", " else:\n", " patience_counter += 1\n", " \n", " if patience_counter >= PATIENCE:\n", " print(f\"ā³ Early stopping triggered at epoch {epoch+1}. Val RMSE has not improved for {PATIENCE} epochs.\")\n", " break\n", " \n", " print(f\"Fold {fold_num} training complete.\") \n", " # --- Final Evaluation on Test Set ---\n", " print(\"Loading best model (based on lowest train loss) for final TEST evaluation...\")\n", " model.load_state_dict(torch.load(fold_model_path)) \n", " \n", " final_test_metrics, y_true_fold, y_pred_fold, pdb_groups_fold, subgroups_fold, sources_fold = evaluate_model(model, test_loader, device, PKD_BOUNDS)\n", " print(f\"Fold {fold_num} Final Test Metrics:\")\n", " for key, value in final_test_metrics.items():\n", " print(f\" - {key.upper()}: {value:.4f}\")\n", "\n", " fold_results.append(final_test_metrics)\n", " all_folds_y_true.extend(y_true_fold)\n", " all_folds_y_pred.extend(y_pred_fold)\n", "\n", " # --- Save Predictions for the Fold ---\n", " # MODIFIED: Create a more detailed DataFrame\n", " pred_df = pd.DataFrame({\n", " 'label': y_true_fold, \n", " 'prediction': y_pred_fold, \n", " 'residual': y_true_fold - y_pred_fold,\n", " 'abs_residual': np.abs(y_true_fold - y_pred_fold),\n", " 'PDB': pdb_groups_fold, \n", " 'Subgroup': subgroups_fold,\n", " 'Source Data Set': sources_fold # Preserved column\n", " })\n", " # MODIFIED: Save to the RESULTS_DIR\n", " pred_filename = os.path.join(RESULTS_DIR, f\"fold_{fold_num}_predictions.csv\")\n", " pred_df.to_csv(pred_filename, index=False)\n", " print(f\" šŸ’¾ Predictions for fold {fold_num} saved to '{pred_filename}'\")\n", "\n", "\n", " # --- Aggregate and Display Final Results ---\n", " print(f\"\\nšŸ“ˆ Generating combined regression plot for {CV_STRATEGY}\")\n", " combined_y_true = np.array(all_folds_y_true)\n", " combined_y_pred = np.array(all_folds_y_pred)\n", "\n", " overall_metrics = {\n", " 'rmse': np.sqrt(mean_squared_error(combined_y_true, combined_y_pred)),\n", " 'pearson': pearsonr(combined_y_true, combined_y_pred)[0],\n", " 'spearman': spearmanr(combined_y_true, combined_y_pred)[0],\n", " 'ci': concordance_index(combined_y_true, combined_y_pred)\n", " }\n", " \n", " plot_regression(\n", " combined_y_true,\n", " combined_y_pred,\n", " overall_metrics,\n", " title=f\"Overall Regression for {CV_STRATEGY} ({N_SPLITS}-Fold CV) with LoRA Fine-tuning\",\n", " filename=os.path.join(RESULTS_DIR, \"overall_regression_plot.png\")\n", " )\n", "\n", " print(f\"\\nāœ… Final Results Summary for {CV_STRATEGY} with LoRA Fine-tuning:\")\n", " results_df = pd.DataFrame(fold_results)\n", " summary_df = pd.DataFrame({'Mean': results_df.mean(), 'Std Dev': results_df.std()})\n", " print(summary_df)\n", "\n", " # NEW: Save the summary metrics to a CSV file\n", " summary_filename = os.path.join(RESULTS_DIR, \"cv_summary_metrics.csv\")\n", " summary_df.to_csv(summary_filename)\n", " print(f\"Overall CV summary metrics saved to {summary_filename}\")\n", " print(f\"Fold model checkpoints and prediction files are saved in the '{RESULTS_DIR}' directory.\")\n", "\n", "# --- Call the new CV function ---\n", "run_sequence_similarity_cv_with_lora()\n", "\n", "\n", "# ==============================================================================\n", "# 8. FINAL NOTES AND RECOMMENDATIONS\n", "# ==============================================================================\n", "print(\"\\n\" + \"=\"*80)\n", "print(\"... 8. FINAL NOTES AND RECOMMENDATIONS ...\")\n", "print(\"\"\"\n", "šŸŽ‰ CONGRATULATIONS! You've successfully implemented BALM with on-the-fly ESM-2 embeddings\n", "and **Parameter-Efficient Fine-Tuning (PEFT) using LoRA**!\n", "\n", "šŸ“Š KEY ADVANTAGES OF THIS APPROACH:\n", "ā€¢ **Efficiency:** Fine-tunes only a tiny fraction of the ESM-2 model parameters (<1%).\n", "ā€¢ **Speed & Memory:** Significantly faster and less memory-intensive than full fine-tuning.\n", "ā€¢ **Performance:** Adapts the powerful ESM-2 model to your specific data, often leading to better performance.\n", "ā€¢ **No Static Cache:** By computing embeddings dynamically, the model fully leverages the fine-tuned LoRA adapters during training.\n", "\n", "šŸš€ PERFORMANCE OPTIMIZATION TIPS:\n", "1. **LoRA Hyperparameters:**\n", " * `lora_rank (r)`: Controls the capacity of the adapters. Higher `r` means more parameters. Typical values are 8-64.\n", " * `lora_alpha`: Acts as a scaling factor. A common heuristic is to set `lora_alpha` to be twice the `lora_rank`. [7]\n", " * `target_modules`: Experiment with adding other modules like `dense` to the target list in the `LoraConfig`. [1]\n", "2. **Learning Rate:** LoRA can sometimes handle slightly higher learning rates than full fine-tuning. Experiment with values from 5e-5 to 5e-4.\n", "3. **Batch Size & Gradient Accumulation:** If your GPU memory is a constraint, you can simulate a larger batch size by using gradient accumulation.\n", "\n", "šŸ’¾ SAVING AND LOADING THE FINE-TUNED MODEL:\n", "ā€¢ After training, you can save the trained LoRA adapters separately from the base model. This is very efficient for storage. [3]\n", "ā€¢ To save: `model.esm_model.save_pretrained(\"path/to/my_lora_adapters\")` [3]\n", "ā€¢ To load: First, load the base ESM-2 model. Then, use `PeftModel.from_pretrained(base_model, \"path/to/my_lora_adapters\")`.\n", "\n", "šŸŽÆ NEXT STEPS:\n", "1. **Hyperparameter Search:** Systematically tune LoRA parameters (`r`, `lora_alpha`) and the learning rate for optimal performance.\n", "2. **Experiment with Larger ESM-2 Models:** LoRA makes fine-tuning even the largest ESM-2 models (e.g., 3B parameters) feasible on a single high-end GPU.\n", "3. **Advanced PEFT Techniques:** Explore other PEFT methods like IA3 or full parameter-efficient fine-tuning for comparison.\n", "\"\"\")\n", "print(\"\\nāœ… LoRA PEFT implementation complete! Happy parameter-efficient protein modeling! šŸ§¬šŸ¤–\")" ] }, { "cell_type": "code", "execution_count": null, "id": "b780f785", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "gpu_env", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.13" } }, "nbformat": 4, "nbformat_minor": 5 }