ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2025-03-21 17:13:14 -07:00
commit 0a750da927
1 changed files with 2 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/models/llama.py
View file

@ -99,11 +99,13 @@ SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__
# Fix new HF's inference code
def _fast_prepare_inputs_for_generation(self, input_ids, **kwargs,):
print("PREPARE", input_ids)
if "past_key_values" in kwargs:
input_ids = input_ids[:,[-1]]
kwargs["attention_mask"] = kwargs["attention_mask"][:,[-1]]
if "cache_position" in kwargs:
kwargs["position_ids"] = kwargs["cache_position"]
print("PREPARE", input_ids)
return { "input_ids" : input_ids, **kwargs, }
pass