Update app.py

app.py

@@ -73,6 +73,14 @@ def compute_loss(model, inputs):
     loss = loss_fct(active_logits, active_labels)
     return loss
 
+# 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):
+        outputs = model(**inputs)
+        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):