From dbf11f7399bc6788f01da12ec5b9eff40a88a845 Mon Sep 17 00:00:00 2001 From: Alvaro Moran Date: Wed, 10 Apr 2024 09:10:07 +0000 Subject: [PATCH] fix: get_generation_mode is now a method of generation_config API change when transformers was updated. --- optimum/tpu/generation/token_selector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/tpu/generation/token_selector.py b/optimum/tpu/generation/token_selector.py index 9cc0994e..e7cef097 100644 --- a/optimum/tpu/generation/token_selector.py +++ b/optimum/tpu/generation/token_selector.py @@ -138,7 +138,7 @@ def create( logger.warning(f"Setting `pad_token_id` to `eos_token_id`:{eos_token_id} for open-end generation.") generation_config.pad_token_id = eos_token_id - generation_mode = model._get_generation_mode(generation_config, None) + generation_mode = model.generation_config.get_generation_mode() if generation_mode not in [GenerationMode.GREEDY_SEARCH, GenerationMode.SAMPLE]: raise ValueError("Unsupported generation mode")