Update test

huggingface · Sep 20, 2024 · 4f51ad4 · 4f51ad4
1 parent f7635ee
commit 4f51ad4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
@@ -60,7 +60,7 @@
         ImageGPTForCausalImageModeling,
         SpeechEncoderDecoderModel,
     )
-    from transformers.cache_utils import DynamicCache, EncoderDecoderCache, QuantoQuantizedCache, StaticCache
+    from transformers.cache_utils import DynamicCache, DynamicSlidingWindowCache, EncoderDecoderCache, QuantoQuantizedCache, StaticCache
     from transformers.generation import (
         BeamSampleDecoderOnlyOutput,
         BeamSampleEncoderDecoderOutput,
@@ -1840,6 +1840,9 @@ def test_new_cache_format(self, num_beams, do_sample):
             if config.is_encoder_decoder:
                 cache_cls = EncoderDecoderCache
                 past_key_values = cache_cls(DynamicCache(), DynamicCache())
+            elif getattr(self.config, "sliding_window", None) is not None:
+                cache_cls = DynamicSlidingWindowCache
+                past_key_values = cache_cls(self.config.sliding_window)
             else:
                 cache_cls = DynamicCache
                 past_key_values = cache_cls()