From 893aab0e573b6f77c7c9f1a75fa6f1264310ca81 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 28 Jan 2024 17:20:49 +1100 Subject: [PATCH] Update dpo.py --- unsloth/models/dpo.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/unsloth/models/dpo.py b/unsloth/models/dpo.py index 92fde81fd0..3ae4d636f4 100644 --- a/unsloth/models/dpo.py +++ b/unsloth/models/dpo.py @@ -101,10 +101,13 @@ pass def PatchDPOTrainer(): - # Patch DPO notebook printing - NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line - from transformers.trainer import DEFAULT_PROGRESS_CALLBACK - DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin - DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log + from transformers.trainer import is_in_notebook + if is_in_notebook(): + # Patch DPO notebook printing + NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line + from transformers.trainer import DEFAULT_PROGRESS_CALLBACK + DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin + DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log + pass pass