From 254988c4044f9bcbe5980e9af5686d40edb453aa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 21 Mar 2025 15:39:43 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 4feee54f35..0bc952a2b8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1017,7 +1017,6 @@ def CausalLM_fast_forward(fast_forward_inference): logits_to_keep: Optional[int] = 0, *args, **kwargs, ) -> Union[Tuple, CausalLMOutputWithPast]: - print(past_key_values) if past_key_values is not None: outputs = fast_forward_inference( self, @@ -2664,6 +2663,13 @@ class FastLlamaModel: model.load_lora = functools.partial(load_lora, model) pass + # Patch generate + if model.generate.__name__ != "unsloth_fast_generate": + model._old_generate = model.generate + unsloth_fast_generate.__doc__ = model._old_generate.__doc__ + model.generate = types.MethodType(unsloth_fast_generate, model) + pass + # Add for_inference and for_training model.for_training = functools.partial(FastLlamaModel.for_training, model) model.for_inference = functools.partial(FastLlamaModel.for_inference, model)