Upload metrics.py

metrics.py

@@ -0,0 +1,116 @@
+import pickle
+import numpy as np
+import torch
+from torch.utils.data import Dataset as TorchDataset
+from transformers import AutoTokenizer, AutoModelForTokenClassification, Trainer
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, roc_auc_score, matthews_corrcoef
+from peft import PeftModel, get_peft_config, PeftConfig, get_peft_model, LoraConfig, TaskType
+from accelerate import Accelerator
+from tqdm import tqdm
+
+# Initialize the Accelerator
+accelerator = Accelerator()
+
+class ProteinDataset(TorchDataset):
+    def __init__(self, sequences_path, labels_path, tokenizer, max_length):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+        with open(sequences_path, "rb") as f:
+            self.sequences = pickle.load(f)
+        
+        with open(labels_path, "rb") as f:
+            self.labels = pickle.load(f)
+
+    def __len__(self):
+        return len(self.sequences)
+
+    def __getitem__(self, idx):
+        sequence = self.sequences[idx]
+        label = self.labels[idx]
+    
+        tokenized = self.tokenizer(sequence, padding='max_length', truncation=True, max_length=self.max_length, return_tensors="pt", is_split_into_words=False, add_special_tokens=False)
+    
+        # Remove the extra batch dimension
+        for key in tokenized:
+            tokenized[key] = tokenized[key].squeeze(0)
+    
+        # Ensure labels are also padded/truncated to match tokenized input
+        label_padded = [-100] * self.max_length  # Using -100 as the ignore index
+        label_padded[:len(label)] = label[:self.max_length]
+        
+        tokenized["labels"] = torch.tensor(label_padded)
+    
+        return tokenized
+
+def compute_metrics(p):
+    predictions, labels = p.predictions, p.label_ids
+    predictions = np.argmax(predictions, axis=2)
+    
+    mask = labels != -100
+    predictions = predictions[mask].flatten()
+    labels = labels[mask].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 evaluate_in_chunks(dataset, trainer, chunk_percentage=0.2):
+    chunk_size = int(len(dataset) * chunk_percentage)
+    all_results = []
+    
+    # Wrap the loop with tqdm for progress bar
+    for i in tqdm(range(0, len(dataset), chunk_size), desc="Evaluating chunks"):
+        chunk = [dataset[j] for j in range(i, min(i + chunk_size, len(dataset)))]
+        chunk_results = trainer.evaluate(chunk)
+        print(f"Results for chunk starting at index {i}: {chunk_results}")
+        
+        # Save the chunk results to disk
+        with open(f"results_chunk_{i}.pkl", "wb") as f:
+            pickle.dump(chunk_results, f)
+        
+        all_results.append(chunk_results)
+
+    return all_results
+
+def aggregate_results(results_list):
+    total_samples = sum([res["eval_samples"] for res in results_list])
+    aggregated_results = {}
+
+    for key in results_list[0].keys():
+        if key == "eval_samples":
+            continue
+        aggregated_results[key] = sum([res[key] * res["eval_samples"] for res in results_list]) / total_samples
+
+    return aggregated_results
+
+# Initialize tokenizer and datasets
+tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D")
+train_dataset = ProteinDataset("data/12M_data/512_train_sequences_chunked_by_family.pkl", "data/12M_data/512_train_labels_chunked_by_family.pkl", tokenizer, 512)
+test_dataset = ProteinDataset("data/12M_data/512_test_sequences_chunked_by_family.pkl", "data/12M_data/512_test_labels_chunked_by_family.pkl", tokenizer, 512)
+
+# Load the pre-trained LoRA model
+base_model_path = "facebook/esm2_t33_650M_UR50D"
+lora_model_path = "qlora_binding_sites/best_model_esm2_t33_650M_qlora_2023-10-18_02-14-48"
+base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
+model = PeftModel.from_pretrained(base_model, lora_model_path)
+model = accelerator.prepare(model)
+
+# Initialize the Trainer
+trainer = Trainer(
+    model=model,
+    compute_metrics=compute_metrics
+)
+
+Evaluate the model on chunks of the training dataset
+train_results = evaluate_in_chunks(train_dataset, trainer)
+aggregated_train_results = aggregate_results(train_results)
+print(f"Aggregated Training Results: {aggregated_train_results}")
+
+# Evaluate the model on chunks of the test dataset
+test_results = evaluate_in_chunks(test_dataset, trainer)
+aggregated_test_results = aggregate_results(test_results)
+print(f"Aggregated Test Results: {aggregated_test_results}")