Update llama.py

This commit is contained in:
Daniel Han 2025-03-21 15:39:43 -07:00
commit 254988c404

View file

@ -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)