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, class_weights): #compute_loss(model, inputs):  add class_weights as input, jw 20240628
+def compute_loss(model, inputs): 
     """Custom compute_loss function."""
     logits = model(**inputs).logits
     labels = inputs["labels"]
@@ -76,11 +76,12 @@ def compute_loss(model, inputs, class_weights): #compute_loss(model, inputs):  a
 # 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, class_weights, return_outputs=False):  #add class_weights as input, jw 20240628
+    def compute_loss(self, model, inputs, return_outputs=False): 
         outputs = model(**inputs)
-        loss = compute_loss(model, inputs, class_weights)  #add class_weights as input, jw 20240628
+        loss = compute_loss(model, inputs) 
         return (loss, outputs) if return_outputs else loss
-        
+
+#         
 # fine-tuning function
 def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset):
     
@@ -196,8 +197,7 @@ def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset)
         eval_dataset=test_dataset,
         tokenizer=tokenizer,
         data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer),
-        compute_metrics=compute_metrics,
-        class_weights=class_weights,  #add class_weights as input, jw 20240628
+        compute_metrics=compute_metrics
     )
 
     # Train and Save Model