Update app.py

app.py

@@ -141,6 +141,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
     class_weights = compute_class_weight(class_weight='balanced', classes=classes, y=flat_train_labels)
     accelerator = Accelerator()
     class_weights = torch.tensor(class_weights, dtype=torch.float32).to(accelerator.device)  
+    print(" class_weights:", class_weights)
     
     # Convert the model into a PeftModel
     peft_config = LoraConfig(
@@ -215,7 +216,7 @@ MODEL_OPTIONS = [
     "facebook/esm2_t33_650M_UR50D",
 ]  # models users can choose from
 
-'''
+
 # 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)
@@ -251,6 +252,7 @@ class_weights = compute_class_weight(class_weight='balanced', classes=classes, y
 accelerator = Accelerator()
 class_weights = torch.tensor(class_weights, dtype=torch.float32).to(accelerator.device)
 
+'''
 # inference
 # Path to the saved LoRA model
 model_path = "AmelieSchreiber/esm2_t12_35M_lora_binding_sites_v2_cp3"