predict loss only in evaluation

SeanLee97 · Sep 29, 2024 · daf2bb6 · daf2bb6
1 parent cfa4f52
commit daf2bb6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/angle_emb/angle.py b/angle_emb/angle.py
@@ -1526,7 +1526,6 @@ def fit(self,
             dataset_format=self.detect_dataset_format(train_ds),
             train_dataset=train_ds,
             eval_dataset=valid_ds,
-            compute_metrics=None,  # do not compute metrics
             loss_kwargs=loss_kwargs,
             tokenizer=self.tokenizer,
             args=TrainingArguments(
@@ -1539,6 +1538,7 @@ def fit(self,
                 logging_steps=logging_steps,
                 save_strategy=save_strategy,
                 evaluation_strategy=evaluation_strategy,
+                prediction_loss_only=True,
                 eval_steps=eval_steps,
                 save_steps=save_steps,
                 output_dir=output_dir,