Update app.py

app.py

@@ -73,6 +73,93 @@ def compute_loss(model, inputs):
     loss = loss_fct(active_logits, active_labels)
     return loss
 
+# fine-tuning function
+def train_function_no_sweeps(base_model_path, train_dataset, test_dataset):
+    
+    # Set the LoRA config
+    config = {
+        "lora_alpha": 1, #try 0.5, 1, 2, ..., 16
+        "lora_dropout": 0.2,
+        "lr": 5.701568055793089e-04,
+        "lr_scheduler_type": "cosine",
+        "max_grad_norm": 0.5,
+        "num_train_epochs": 3,
+        "per_device_train_batch_size": 12,
+        "r": 2,
+        "weight_decay": 0.2,
+        # Add other hyperparameters as needed
+    }
+    # The base model you will train a LoRA on top of
+    base_model_path = "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()}
+    base_model = AutoModelForTokenClassification.from_pretrained(base_model_path, 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 try "dense_h_to_4h" and "dense_4h_to_h"
+        lora_dropout=config["lora_dropout"], 
+        bias="none" # or "all" or "lora_only" 
+    )
+    base_model = get_peft_model(base_model, peft_config)
+
+    # Use the accelerator
+    base_model = accelerator.prepare(base_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=base_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)
+
 # Load the data from pickle files (replace with your local paths)
 with open("./datasets/train_sequences_chunked_by_family.pkl", "rb") as f:
     train_sequences = pickle.load(f)