Upload qlora_train_v2.py

qlora_train_v2.py

@@ -0,0 +1,283 @@
+import os
+import wandb
+import numpy as np
+import torch
+import torch.nn as nn
+from datetime import datetime
+from sklearn.model_selection import train_test_split
+from sklearn.utils.class_weight import compute_class_weight
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score, matthews_corrcoef
+from transformers import (
+    AutoModelForTokenClassification,
+    AutoTokenizer,
+    DataCollatorForTokenClassification,
+    TrainingArguments,
+    Trainer,
+    BitsAndBytesConfig
+)
+from datasets import Dataset
+from accelerate import Accelerator
+from peft import get_peft_config, PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training
+import pickle
+import gc
+from tqdm import tqdm
+
+# Initialize accelerator and Weights & Biases
+accelerator = Accelerator()
+os.environ["WANDB_NOTEBOOK_NAME"] = 'qlora_train.py'
+wandb.init(project='binding_site_prediction')
+
+# Helper Functions and Data Preparation
+#-----------------------------------------------------------------------------
+# Added this first function in
+def print_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    print(
+        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+    )
+#-----------------------------------------------------------------------------
+
+def save_config_to_txt(config, filename):
+    """Save the configuration dictionary to a text file."""
+    with open(filename, 'w') as f:
+        for key, value in config.items():
+            f.write(f"{key}: {value}\n")
+
+def truncate_labels(labels, max_length):
+    return [label[:max_length] for label in labels]
+
+def compute_metrics(p):
+    predictions, labels = p
+    predictions = np.argmax(predictions, axis=2)
+    predictions = predictions[labels != -100].flatten()
+    labels = labels[labels != -100].flatten()
+    accuracy = accuracy_score(labels, predictions)
+    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')
+    auc = roc_auc_score(labels, predictions)
+    mcc = matthews_corrcoef(labels, predictions)
+    return {'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'auc': auc, 'mcc': mcc}
+
+def compute_loss(model, logits, inputs):
+    # logits = model(**inputs).logits
+    labels = inputs["labels"]
+    loss_fct = nn.CrossEntropyLoss(weight=class_weights)
+    active_loss = inputs["attention_mask"].view(-1) == 1
+    active_logits = logits.view(-1, model.config.num_labels)
+    active_labels = torch.where(
+        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+    )
+    loss = loss_fct(active_logits, active_labels)
+    return loss
+
+# Load data from pickle files
+with open("data/16M_data_small/v3_train_sequences_chunked_by_family.pkl", "rb") as f:
+    train_sequences = pickle.load(f)
+    
+with open("data/16M_data_small/v3_test_sequences_chunked_by_family.pkl", "rb") as f:
+    test_sequences = pickle.load(f)
+
+with open("data/16M_data_small/v3_train_labels_chunked_by_family.pkl", "rb") as f:
+    train_labels = pickle.load(f)
+
+with open("data/16M_data_small/v3_test_labels_chunked_by_family.pkl", "rb") as f:
+    test_labels = pickle.load(f)
+
+
+# Tokenization
+tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t30_150M_UR50D")
+
+# Set max_sequence_length to the tokenizer's max input length
+max_sequence_length = tokenizer.model_max_length
+
+# Function to tokenize in batches, show progress bar with tqdm, and free up memory
+def batch_tokenize(sequences, batch_size=10000):  # Reduced batch size for better memory management
+    tokenized_outputs = {'input_ids': [], 'attention_mask': []}  # Initialize empty lists for outputs
+    
+    for i in tqdm(range(0, len(sequences), batch_size)):
+        batch = sequences[i:i+batch_size]
+        tokenized_batch = tokenizer(batch, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
+        
+        # Append the tokenized batch data to the outputs
+        for key in tokenized_outputs:
+            tokenized_outputs[key].append(tokenized_batch[key])
+        
+        # Free up memory
+        del batch
+        del tokenized_batch
+        gc.collect()
+
+    # Convert lists to tensors
+    for key in tokenized_outputs:
+        tokenized_outputs[key] = torch.cat(tokenized_outputs[key])
+
+    return tokenized_outputs
+
+train_tokenized = batch_tokenize(train_sequences)
+test_tokenized = batch_tokenize(test_sequences)
+
+# Free memory
+del train_sequences
+del test_sequences
+gc.collect()
+
+# Directly truncate the entire list of labels
+train_labels = truncate_labels(train_labels, max_sequence_length)
+test_labels = truncate_labels(test_labels, max_sequence_length)
+
+train_dataset = Dataset.from_dict({k: v for k, v in train_tokenized.items()}).add_column("labels", train_labels)
+test_dataset = Dataset.from_dict({k: v for k, v in test_tokenized.items()}).add_column("labels", test_labels)
+
+# Compute Class Weights
+classes = [0, 1]
+flat_train_labels = [label for sublist in train_labels for label in sublist]
+class_weights = compute_class_weight(class_weight='balanced', classes=classes, y=flat_train_labels)
+class_weights = torch.tensor(class_weights, dtype=torch.bfloat16).to(accelerator.device)
+
+# Define Custom Trainer Class
+class WeightedTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False):
+        outputs = model(**inputs)
+        logits = outputs.logits
+        loss = compute_loss(model, logits, inputs)
+        return (loss, outputs) if return_outputs else loss
+
+# Configure the quantization settings
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+def train_function_no_sweeps(train_dataset, test_dataset):
+    
+    # Directly set the config
+    config = {
+        "lora_alpha": 1, 
+        "lora_dropout": 0.5,
+        "lr": 1.701568055793089e-04,
+        "lr_scheduler_type": "cosine",
+        "max_grad_norm": 0.5,
+        "num_train_epochs": 5,
+        "per_device_train_batch_size": 40,
+        "r": 2,
+        "weight_decay": 0.3,
+        # Add other hyperparameters as needed
+    }
+
+    # Log the config to W&B
+    wandb.config.update(config)
+
+    # Save the config to a text file
+    timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+    config_filename = f"esm2_t30_150M_qlora_config_{timestamp}.txt"
+    save_config_to_txt(config, config_filename)
+          
+    model_checkpoint = "facebook/esm2_t30_150M_UR50D"  
+    
+    # Define labels and model
+    id2label = {0: "No binding site", 1: "Binding site"}
+    label2id = {v: k for k, v in id2label.items()}
+    
+    model = AutoModelForTokenClassification.from_pretrained(
+        model_checkpoint,
+        num_labels=len(id2label),
+        id2label=id2label,
+        label2id=label2id,
+        quantization_config=bnb_config  # Apply quantization here
+    )
+
+    # Prepare the model for 4-bit quantization training
+    model.gradient_checkpointing_enable()
+    model = prepare_model_for_kbit_training(model)
+    
+    # Convert the model into a PeftModel
+    peft_config = LoraConfig(
+        task_type=TaskType.TOKEN_CLS,
+        inference_mode=False,
+        r=config["r"],
+        lora_alpha=config["lora_alpha"],
+        target_modules=[
+            "query",
+            "key",
+            "value",
+            "EsmSelfOutput.dense",
+            "EsmIntermediate.dense",
+            "EsmOutput.dense",
+            "EsmContactPredictionHead.regression",
+            "classifier"
+        ],
+        lora_dropout=config["lora_dropout"],
+        bias="none",  # or "all" or "lora_only"
+        # modules_to_save=["classifier"]
+    )
+    model = get_peft_model(model, peft_config)
+    print_trainable_parameters(model) # added this in
+
+    # Use the accelerator
+    model = accelerator.prepare(model)
+    train_dataset = accelerator.prepare(train_dataset)
+    test_dataset = accelerator.prepare(test_dataset)
+
+    timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+
+    # Training setup
+    training_args = TrainingArguments(
+        output_dir=f"esm2_t30_150M_qlora_binding_sites_{timestamp}",
+        learning_rate=config["lr"],
+        lr_scheduler_type=config["lr_scheduler_type"],
+        gradient_accumulation_steps=4, # changed from 1 to 4
+        # warmup_steps=2, # added this in 
+        max_grad_norm=config["max_grad_norm"],
+        per_device_train_batch_size=config["per_device_train_batch_size"],
+        per_device_eval_batch_size=config["per_device_train_batch_size"],
+        num_train_epochs=config["num_train_epochs"],
+        weight_decay=config["weight_decay"],
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        metric_for_best_model="f1",
+        greater_is_better=True,
+        push_to_hub=False,
+        logging_dir=None,
+        logging_first_step=False,
+        logging_steps=200,
+        save_total_limit=7,
+        no_cuda=False,
+        seed=8893,
+        fp16=True,
+        report_to='wandb', 
+        optim="paged_adamw_8bit" # added this in 
+
+    )
+    
+    # Initialize Trainer
+    trainer = WeightedTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=test_dataset,
+        tokenizer=tokenizer,
+        data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer),
+        compute_metrics=compute_metrics
+    )
+
+    # Train and Save Model
+    trainer.train()
+    save_path = os.path.join("qlora_binding_sites", f"best_model_esm2_t30_150M_qlora_{timestamp}")
+    trainer.save_model(save_path)
+    tokenizer.save_pretrained(save_path)
+
+# Call the training function
+if __name__ == "__main__":
+    train_function_no_sweeps(train_dataset, test_dataset)
+
+