Upload 3 files

ensemble (1).py

@@ -0,0 +1,113 @@
+import os
+import pickle
+import numpy as np
+from scipy import stats
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score, matthews_corrcoef
+from transformers import AutoModelForTokenClassification, Trainer, AutoTokenizer, DataCollatorForTokenClassification
+from datasets import Dataset, concatenate_datasets
+from accelerate import Accelerator
+from peft import PeftModel
+import gc
+
+# Step 1: Load train/test data and labels from pickle files
+with open("/kaggle/input/550k-dataset/train_sequences_chunked_by_family.pkl", "rb") as f:
+    train_sequences = pickle.load(f)
+with open("/kaggle/input/550k-dataset/test_sequences_chunked_by_family.pkl", "rb") as f:
+    test_sequences = pickle.load(f)
+with open("/kaggle/input/550k-dataset/train_labels_chunked_by_family.pkl", "rb") as f:
+    train_labels = pickle.load(f)
+with open("/kaggle/input/550k-dataset/test_labels_chunked_by_family.pkl", "rb") as f:
+    test_labels = pickle.load(f)
+
+# Step 2: Define the Tokenizer
+tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
+max_sequence_length = tokenizer.model_max_length
+
+# Step 3: Define a `compute_metrics_for_batch` function.
+def compute_metrics_for_batch(sequences_batch, labels_batch, models, voting='hard'):
+    # Tokenize batch
+    batch_tokenized = tokenizer(sequences_batch, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
+    # print("Shape of tokenized sequences:", batch_tokenized["input_ids"].shape)  # Debug print
+    
+    batch_dataset = Dataset.from_dict({k: v for k, v in batch_tokenized.items()})
+    batch_dataset = batch_dataset.add_column("labels", labels_batch[:len(batch_dataset)])
+    
+    # Convert labels to numpy array of shape (1000, 1002)
+    labels_array = np.array([np.pad(label, (0, 1002 - len(label)), constant_values=-100) for label in batch_dataset["labels"]])
+    
+    # Initialize a trainer for each model
+    data_collator = DataCollatorForTokenClassification(tokenizer)
+    trainers = [Trainer(model=model, data_collator=data_collator) for model in models]
+    
+    # Get the predictions from each model
+    all_predictions = [trainer.predict(test_dataset=batch_dataset)[0] for trainer in trainers]
+    
+    if voting == 'hard':
+        # Hard voting
+        hard_predictions = [np.argmax(predictions, axis=2) for predictions in all_predictions]
+        ensemble_predictions = stats.mode(hard_predictions, axis=0)[0][0]
+    elif voting == 'soft':
+        # Soft voting
+        avg_predictions = np.mean(all_predictions, axis=0)
+        ensemble_predictions = np.argmax(avg_predictions, axis=2)
+    else:
+        raise ValueError("Voting must be either 'hard' or 'soft'")
+        
+    # Use broadcasting to create 2D mask
+    mask_2d = labels_array != -100
+    
+    # Filter true labels and predictions using the mask
+    true_labels_list = [label[mask_2d[idx]] for idx, label in enumerate(labels_array)]
+    true_labels = np.concatenate(true_labels_list)
+    flat_predictions_list = [ensemble_predictions[idx][mask_2d[idx]] for idx in range(ensemble_predictions.shape[0])]
+    flat_predictions = np.concatenate(flat_predictions_list).tolist()
+
+    # Compute the metrics
+    accuracy = accuracy_score(true_labels, flat_predictions)
+    precision, recall, f1, _ = precision_recall_fscore_support(true_labels, flat_predictions, average='binary')
+    auc = roc_auc_score(true_labels, flat_predictions)
+    mcc = matthews_corrcoef(true_labels, flat_predictions)  # Compute MCC
+    
+    return {"accuracy": accuracy, "precision": precision, "recall": recall, "f1": f1, "auc": auc, "mcc": mcc}
+
+# Step 4: Evaluate in batches
+def evaluate_in_batches(sequences, labels, models, dataset_name, voting, batch_size=1000, print_first_n=5):
+    num_batches = len(sequences) // batch_size + int(len(sequences) % batch_size != 0)
+    metrics_list = []
+    
+    for i in range(num_batches):
+        start_idx = i * batch_size
+        end_idx = start_idx + batch_size
+        batch_metrics = compute_metrics_for_batch(sequences[start_idx:end_idx], labels[start_idx:end_idx], models, voting)
+        
+        # Print metrics for the first few batches for both train and test datasets
+        if i < print_first_n:
+            print(f"{dataset_name} - Batch {i+1}/{num_batches} metrics: {batch_metrics}")
+        
+        metrics_list.append(batch_metrics)
+    
+    # Average metrics over all batches
+    avg_metrics = {key: np.mean([metrics[key] for metrics in metrics_list]) for key in metrics_list[0]}
+    return avg_metrics
+
+# Step 5: Load pre-trained base model and fine-tuned LoRA models
+accelerator = Accelerator()
+base_model_path = "facebook/esm2_t12_35M_UR50D"
+base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
+lora_model_paths = [
+    "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_cp1",
+    "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp1",
+]
+models = [PeftModel.from_pretrained(base_model, path) for path in lora_model_paths]
+models = [accelerator.prepare(model) for model in models]
+
+# Step 6: Compute and print the metrics
+test_metrics_soft = evaluate_in_batches(test_sequences, test_labels, models, "test", voting='soft')
+train_metrics_soft = evaluate_in_batches(train_sequences, train_labels, models, "train", voting='soft')
+test_metrics_hard = evaluate_in_batches(test_sequences, test_labels, models, "test", voting='hard')
+train_metrics_hard = evaluate_in_batches(train_sequences, train_labels, models, "train", voting='hard')
+
+print("Test metrics (soft voting):", test_metrics_soft)
+print("Train metrics (soft voting):", train_metrics_soft)
+print("Test metrics (hard voting):", test_metrics_hard)
+print("Train metrics (hard voting):", train_metrics_hard)

metrics (1).py

@@ -0,0 +1,95 @@
+import os
+import wandb
+import numpy as np
+import pickle
+import torch
+import torch.nn as nn
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score, matthews_corrcoef
+from transformers import AutoModelForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification, Trainer
+from datasets import Dataset
+from accelerate import Accelerator
+from peft import PeftModel
+
+# Helper functions and data preparation
+def truncate_labels(labels, max_length):
+    """Truncate labels to the specified max_length."""
+    return [label[:max_length] for label in labels]
+
+def compute_metrics(p):
+    """Compute metrics for evaluation."""
+    predictions, labels = p
+    predictions = np.argmax(predictions, axis=2)
+    
+    # Remove padding (-100 labels)
+    predictions = predictions[labels != -100].flatten()
+    labels = labels[labels != -100].flatten()
+    
+    # Compute accuracy
+    accuracy = accuracy_score(labels, predictions)
+    
+    # Compute precision, recall, F1 score, and AUC
+    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average='binary')
+    auc = roc_auc_score(labels, predictions)
+    
+    # Compute MCC
+    mcc = matthews_corrcoef(labels, predictions)
+    
+    return {'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'auc': auc, 'mcc': mcc}
+
+class WeightedTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False):
+        """Custom compute_loss function."""
+        outputs = model(**inputs)
+        loss_fct = nn.CrossEntropyLoss()
+        active_loss = inputs["attention_mask"].view(-1) == 1
+        active_logits = outputs.logits.view(-1, model.config.num_labels)
+        active_labels = torch.where(
+            active_loss, inputs["labels"].view(-1), torch.tensor(loss_fct.ignore_index).type_as(inputs["labels"])
+        )
+        loss = loss_fct(active_logits, active_labels)
+        return (loss, outputs) if return_outputs else loss
+
+if __name__ == "__main__":
+    # Environment setup
+    accelerator = Accelerator()
+    wandb.init(project='binding_site_prediction')
+    
+    # Load data and labels
+    with open("600K_data/train_sequences_chunked_by_family.pkl", "rb") as f:
+        train_sequences = pickle.load(f)
+    with open("600K_data/test_sequences_chunked_by_family.pkl", "rb") as f:
+        test_sequences = pickle.load(f)
+    with open("600K_data/train_labels_chunked_by_family.pkl", "rb") as f:
+        train_labels = pickle.load(f)
+    with open("600K_data/test_labels_chunked_by_family.pkl", "rb") as f:
+        test_labels = pickle.load(f)
+
+    # Tokenization and dataset creation
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
+    max_sequence_length = tokenizer.model_max_length
+    train_tokenized = tokenizer(train_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
+    test_tokenized = tokenizer(test_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
+    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)
+
+    # Load the pre-trained LoRA model
+    base_model_path = "facebook/esm2_t12_35M_UR50D"
+    lora_model_path = "esm2_t12_35M_lora_binding_sites_2023-09-21_17-50-58/checkpoint-84029" # Replace with the correct path to your LoRA model
+    base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
+    model = PeftModel.from_pretrained(base_model, lora_model_path)
+    model = accelerator.prepare(model)
+
+    # Define a function to compute metrics and get the train/test metrics
+    data_collator = DataCollatorForTokenClassification(tokenizer)
+    trainer = Trainer(model=model, data_collator=data_collator, compute_metrics=compute_metrics)
+    train_metrics = trainer.evaluate(train_dataset)
+    test_metrics = trainer.evaluate(test_dataset)
+    
+    # Print the metrics
+    print(f"Train metrics: {train_metrics}")
+    print(f"Test metrics: {test_metrics}")
+
+    # Log metrics to W&B
+    wandb.log({"Train metrics": train_metrics, "Test metrics": test_metrics})

train (1).py

@@ -0,0 +1,193 @@
+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
+from datasets import Dataset
+from accelerate import Accelerator
+from peft import get_peft_config, PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType
+import pickle
+
+# Initialize accelerator and Weights & Biases
+accelerator = Accelerator()
+os.environ["WANDB_NOTEBOOK_NAME"] = 'train.py'
+wandb.init(project='binding_site_prediction')
+
+# Helper Functions and Data Preparation
+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, 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("600K_data/train_sequences_chunked_by_family.pkl", "rb") as f:
+    train_sequences = pickle.load(f)
+    
+with open("600K_data/test_sequences_chunked_by_family.pkl", "rb") as f:
+    test_sequences = pickle.load(f)
+
+with open("600K_data/train_labels_chunked_by_family.pkl", "rb") as f:
+    train_labels = pickle.load(f)
+
+with open("600K_data/test_labels_chunked_by_family.pkl", "rb") as f:
+    test_labels = pickle.load(f)
+    
+# Tokenization
+tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
+
+# Set max_sequence_length to the tokenizer's max input length
+max_sequence_length = tokenizer.model_max_length
+
+train_tokenized = tokenizer(train_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
+test_tokenized = tokenizer(test_sequences, padding=True, truncation=True, max_length=max_sequence_length, return_tensors="pt", is_split_into_words=False)
+
+# 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.float32).to(accelerator.device)
+
+# Define Custom Trainer Class
+class WeightedTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False):
+        outputs = model(**inputs)
+        loss = compute_loss(model, inputs)
+        return (loss, outputs) if return_outputs else loss
+
+# Define and run training function
+def train_function_no_sweeps(train_dataset, test_dataset):
+    
+    # Directly set the config
+    config = {
+        "lora_alpha": 1, 
+        "lora_dropout": 0.4,
+        "lr": 5.701568055793089e-04,
+        "lr_scheduler_type": "cosine",
+        "max_grad_norm": 0.5,
+        "num_train_epochs": 1,
+        "per_device_train_batch_size": 6,
+        "r": 1,
+        "weight_decay": 0.4,
+        # 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_t12_35M_lora_config_{timestamp}.txt"
+    save_config_to_txt(config, config_filename)
+    
+    model_checkpoint = "facebook/esm2_t12_35M_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)
+
+    # 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"], # also maybe "dense_h_to_4h" and "dense_4h_to_h"
+        lora_dropout=config["lora_dropout"], 
+        bias="none" # or "all" or "lora_only" 
+    )
+    model = get_peft_model(model, peft_config)
+
+    # 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_t12_35M_lora_binding_sites_{timestamp}",
+        learning_rate=config["lr"],
+        lr_scheduler_type=config["lr_scheduler_type"],
+        gradient_accumulation_steps=1,
+        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'
+    )
+
+    # 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("lora_binding_sites", f"best_model_esm2_t12_35M_lora_{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)