From 81b1cc3a22038ba0df0b7fa0d1e6d4a127f1a48a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 13 Feb 2025 16:27:02 -0800 Subject: [PATCH] Update rl_replacements.py --- unsloth/models/rl_replacements.py | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 4d7a4dbe09..0a6ea5dff4 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -101,23 +101,20 @@ RL_FUNCTIONS["sft_trainer"].append(sft_trainer_prepare_dataset) # Ignore mean_token_accuracy since it needs logits # We override it directly with our version -def _sft_trainer_compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None): - (loss, outputs) = super().compute_loss( - model, - inputs, - return_outputs = return_outputs, - num_items_in_batch = num_items_in_batch, - ) - return (loss, outputs) if return_outputs else loss -pass - def sft_trainer_compute_loss(function_name, function): if function_name != "compute_loss": return function - function = inspect.getsource(_sft_trainer_compute_loss) - function = function.replace("def _sft_trainer_compute_loss", "def compute_loss") - function = function.split("\n") - function = "\n".join(" "*4+x for x in function) + def compute_loss(self, model, inputs, return_outputs = False, num_items_in_batch = None): + (loss, outputs) = super().compute_loss( + model, + inputs, + return_outputs = return_outputs, + num_items_in_batch = num_items_in_batch, + ) + return (loss, outputs) if return_outputs else loss + pass + + function = inspect.getsource(compute_loss) return function pass RL_FUNCTIONS["sft_trainer"].append(sft_trainer_compute_loss)