Update app.py

app.py

@@ -60,7 +60,7 @@ def compute_metrics(p):
     
     return {'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1': f1, 'auc': auc, 'mcc': mcc} 
 
-def compute_loss(model, inputs):
+def compute_loss(model, inputs, class_weights): #compute_loss(model, inputs):  add class_weights as input, jw 20240628
     """Custom compute_loss function."""
     logits = model(**inputs).logits
     labels = inputs["labels"]
@@ -76,9 +76,9 @@ def compute_loss(model, inputs):
 # Define Custom Trainer Class
 # Since we are using class weights, due to the imbalance between non-binding residues and binding residues, we will need a custom weighted trainer.
 class WeightedTrainer(Trainer):
-    def compute_loss(self, model, inputs, return_outputs=False):
+    def compute_loss(self, model, inputs, class_weights, return_outputs=False):  #add class_weights as input, jw 20240628
         outputs = model(**inputs)
-        loss = compute_loss(model, inputs)
+        loss = compute_loss(model, inputs, class_weights)  #add class_weights as input, jw 20240628
         return (loss, outputs) if return_outputs else loss
         
 # fine-tuning function
@@ -185,7 +185,8 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         seed=8893,
         fp16=True,
         #report_to='wandb'
-        report_to=None
+        report_to=None,
+        class_weights=class_weights,  #jw, 20240628
     )
 
     # Initialize Trainer