From ba515ec92dbc85c03c65d3f31e10166cc73ef323 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 20 Jul 2024 12:47:36 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 32610bbfda..ff51b90b84 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1068,7 +1068,14 @@ def _wrap_fast_inference(generate, device_type, dtype, model): # For newer HF kwargs["cache_implementation"] = "dynamic" - print(kwargs) + # Remove token_type_ids + kwargs.pop("token_type_ids", None) + + # Check pad_token + kwargs["pad_token_id"] = kwargs.pop( + "pad_token_id", + getattr(model.config, "eos_token_id", None), + ) # Set pad token # old_pad_token_id = getattr(model.config, "pad_token_id", None)