From 4e9c09556c0ccf679249e6c680979a606e0ac7e4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 21 Sep 2024 01:56:17 -0700 Subject: [PATCH] Update chat_templates.py --- unsloth/chat_templates.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index e19bea0771..1b59e379a7 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -1841,7 +1841,11 @@ def train_on_responses_only( pass return { "labels" : all_labels } pass - trainer.train_dataset = trainer.train_dataset.map(_train_on_responses_only, batched = True) + + if hasattr(trainer, "train_dataset"): + trainer.train_dataset = trainer.train_dataset.map(_train_on_responses_only, batched = True) + if hasattr(trainer, "eval_dataset"): + trainer.eval_dataset = trainer.eval_dataset .map(_train_on_responses_only, batched = True) return trainer pass