|
import os |
|
import wandb |
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
from datetime import datetime |
|
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, |
|
default_data_collator |
|
) |
|
from torch.utils.data import Dataset as TorchDataset |
|
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 |
|
|
|
|
|
accelerator = Accelerator() |
|
os.environ["WANDB_NOTEBOOK_NAME"] = 'qlora_train.py' |
|
wandb.init(project='binding_site_prediction') |
|
|
|
|
|
|
|
|
|
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): |
|
|
|
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 |
|
|
|
|
|
with open("data/16M_data_big/v2_train_sequences_chunked_by_family.pkl", "rb") as f: |
|
train_sequences = pickle.load(f) |
|
del f |
|
gc.collect() |
|
|
|
with open("data/16M_data_big/v2_test_sequences_chunked_by_family.pkl", "rb") as f: |
|
test_sequences = pickle.load(f) |
|
del f |
|
gc.collect() |
|
|
|
with open("data/16M_data_big/v2_train_labels_chunked_by_family.pkl", "rb") as f: |
|
train_labels = pickle.load(f) |
|
del f |
|
gc.collect() |
|
|
|
with open("data/16M_data_big/v2_test_labels_chunked_by_family.pkl", "rb") as f: |
|
test_labels = pickle.load(f) |
|
del f |
|
gc.collect() |
|
|
|
|
|
desired_length = 1022 |
|
tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D") |
|
sample_sequence = "A" |
|
tokenized_sample = tokenizer(sample_sequence) |
|
|
|
|
|
print(f"Sample Sequence: {sample_sequence}") |
|
print(f"Tokenized Sample: {tokenized_sample}") |
|
print(f"Number of tokens in tokenized sample: {len(tokenized_sample['input_ids'])}") |
|
|
|
num_special_tokens = len(tokenized_sample["input_ids"]) - 1 |
|
print(f"Number of special tokens: {num_special_tokens}") |
|
|
|
effective_length = desired_length - num_special_tokens |
|
print(f"Effective sequence length (accounting for special tokens): {effective_length}") |
|
|
|
|
|
class CustomDataset(TorchDataset): |
|
def __init__(self, sequences, labels, tokenizer, max_length): |
|
self.sequences = sequences |
|
self.labels = labels |
|
self.tokenizer = tokenizer |
|
self.max_length = max_length |
|
|
|
def __len__(self): |
|
return len(self.sequences) |
|
|
|
def __getitem__(self, idx): |
|
sequence = self.sequences[idx] |
|
label = self.labels[idx][:self.max_length] |
|
|
|
tokenized = self.tokenizer(sequence, padding='max_length', truncation=True, max_length=effective_length, return_tensors="pt", is_split_into_words=False) |
|
|
|
|
|
for key, value in tokenized.items(): |
|
tokenized[key] = value[0] |
|
|
|
tokenized['labels'] = torch.tensor(label, dtype=torch.long) |
|
|
|
|
|
|
|
|
|
|
|
del sequence, label |
|
gc.collect() |
|
|
|
return tokenized |
|
|
|
|
|
train_dataset = CustomDataset(train_sequences, train_labels, tokenizer, effective_length) |
|
test_dataset = CustomDataset(test_sequences, test_labels, tokenizer, effective_length) |
|
|
|
|
|
|
|
classes = [0, 1] |
|
|
|
flat_train_labels_gen = (label for sublist in tqdm(train_labels, desc="Flattening labels") for label in sublist) |
|
flat_train_labels = np.fromiter(flat_train_labels_gen, dtype=np.int8) |
|
|
|
del train_sequences, test_sequences, test_labels |
|
gc.collect() |
|
|
|
def compute_average_class_weight(train_labels, classes, batch_size): |
|
num_batches = len(train_labels) // batch_size + (len(train_labels) % batch_size != 0) |
|
total_weights = np.zeros(len(classes)) |
|
|
|
for i in tqdm(range(num_batches), desc="Computing class weights in batches"): |
|
start_idx = i * batch_size |
|
end_idx = start_idx + batch_size |
|
|
|
batch_labels = train_labels[start_idx:end_idx] |
|
flat_labels = np.array([label for sublist in batch_labels for label in sublist], dtype=np.int8) |
|
|
|
weights = compute_class_weight(class_weight='balanced', classes=classes, y=flat_labels) |
|
total_weights += weights |
|
|
|
|
|
del batch_labels, flat_labels, weights |
|
gc.collect() |
|
|
|
|
|
average_weights = total_weights / num_batches |
|
return average_weights |
|
|
|
batch_size = 100000 |
|
class_weights = compute_average_class_weight(train_labels, classes, batch_size) |
|
class_weights = torch.tensor(class_weights, dtype=torch.float32).to(accelerator.device) |
|
|
|
del train_labels |
|
gc.collect() |
|
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
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): |
|
|
|
|
|
config = { |
|
"lora_alpha": 1, |
|
"lora_dropout": 0.5, |
|
"lr": 1.701568055793089e-04, |
|
"lr_scheduler_type": "cosine", |
|
"max_grad_norm": 0.5, |
|
"num_train_epochs": 4, |
|
"per_device_train_batch_size": 60, |
|
"r": 2, |
|
"weight_decay": 0.3, |
|
|
|
} |
|
|
|
|
|
wandb.config.update(config) |
|
|
|
|
|
timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') |
|
config_filename = f"esm2_t33_650M_qlora_config_{timestamp}.txt" |
|
save_config_to_txt(config, config_filename) |
|
|
|
model_checkpoint = "facebook/esm2_t33_650M_UR50D" |
|
|
|
|
|
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 |
|
) |
|
|
|
|
|
model.gradient_checkpointing_enable() |
|
model = prepare_model_for_kbit_training(model) |
|
|
|
|
|
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", |
|
|
|
) |
|
model = get_peft_model(model, peft_config) |
|
print_trainable_parameters(model) |
|
|
|
|
|
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_args = TrainingArguments( |
|
output_dir=f"esm2_t33_650M_qlora_binding_sites_{timestamp}", |
|
learning_rate=config["lr"], |
|
lr_scheduler_type=config["lr_scheduler_type"], |
|
gradient_accumulation_steps=4, |
|
|
|
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" |
|
|
|
) |
|
|
|
|
|
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 |
|
) |
|
|
|
|
|
trainer.train() |
|
save_path = os.path.join("qlora_binding_sites", f"best_model_esm2_t33_650M_qlora_{timestamp}") |
|
trainer.save_model(save_path) |
|
tokenizer.save_pretrained(save_path) |
|
|
|
|
|
if __name__ == "__main__": |
|
train_function_no_sweeps(train_dataset, test_dataset) |
|
|
|
|
|
|
|
|