From 92b3752cad3c4e3bb9f9ee0a08d7d4c4bbffe5b6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 23 Jul 2024 12:27:46 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 338ae0a7ce..719aee5374 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1178,10 +1178,11 @@ def _wrap_fast_inference(generate, device_type, dtype, model): 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), - ) + model_eos_token_id = getattr(model.config, "eos_token_id", None) + if hasattr(model_eos_token_id, "__iter__"): + model_eos_token_id = model_eos_token_id[0] + + kwargs["pad_token_id"] = kwargs.pop("pad_token_id", model_eos_token_id) # Set pad token # old_pad_token_id = getattr(model.config, "pad_token_id", None)