Update app.py

app.py

@@ -196,11 +196,12 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     train_dataset = accelerator.prepare(train_dataset)
     test_dataset = accelerator.prepare(test_dataset)
 
+    model_name_base = base_model_path.split("/")[1]
     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}",
+        output_dir=f"{model_name_base}-lora-binding-sites_{timestamp}",
         learning_rate=config["lr"],
         lr_scheduler_type=config["lr_scheduler_type"],
         gradient_accumulation_steps=1,
@@ -241,9 +242,9 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
 
     # 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)
+    #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)
 
     return save_path