From 5777c3cb3f7a6a82bddeae378ad9001cd796c315 Mon Sep 17 00:00:00 2001 From: Yih-Dar <2521628+ydshieh@users.noreply.github.com> Date: Thu, 18 May 2023 14:54:23 +0200 Subject: [PATCH] Fix (skip) a pipeline test for `RwkvModel` (#23444) fix Co-authored-by: ydshieh --- tests/pipelines/test_pipelines_text_generation.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/pipelines/test_pipelines_text_generation.py b/tests/pipelines/test_pipelines_text_generation.py index 84b1449052..b4605164ae 100644 --- a/tests/pipelines/test_pipelines_text_generation.py +++ b/tests/pipelines/test_pipelines_text_generation.py @@ -240,7 +240,11 @@ class TextGenerationPipelineTests(unittest.TestCase): # We don't care about infinite range models. # They already work. # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly. - if tokenizer.model_max_length < 10000 and "XGLM" not in tokenizer.__class__.__name__: + EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS = ["RwkvForCausalLM", "XGLMForCausalLM"] + if ( + tokenizer.model_max_length < 10000 + and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS + ): # Handling of large generations with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)): text_generator("This is a test" * 500, max_new_tokens=20)