save eval_dataset for subsequent calls (#1125)

huggingface · Dec 21, 2023 · 06b7959 · 06b7959
1 parent b07935f
commit 06b7959
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/trl/trainer/dpo_trainer.py b/trl/trainer/dpo_trainer.py
@@ -492,6 +492,9 @@ def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoa
                 name="reference_rejected_logps", column=all_reference_rejected_logps
             )
 
+            # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs
+            if self.eval_dataset is not None:
+                self.eval_dataset = eval_dataset
             self._precomputed_eval_ref_log_probs = True
 
         return super().get_eval_dataloader(eval_dataset=eval_dataset)