update cache conversion in test

huggingface · Sep 20, 2024 · 4f2e6b2 · 4f2e6b2
1 parent 617726c
commit 4f2e6b2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py
@@ -1907,7 +1907,10 @@ def test_new_cache_format(self, num_beams, do_sample):
                     )
 
             new_cache = new_results.past_key_values
-            legacy_cache_converted = cache_cls.from_legacy_cache(legacy_results.past_key_values)
+            if cache_cls == DynamicSlidingWindowCache:
+                legacy_cache_converted = cache_cls.from_legacy_cache(config.sliding_window, legacy_results.past_key_values)
+            else:
+                legacy_cache_converted = cache_cls.from_legacy_cache(legacy_results.past_key_values)
             for layer_idx in range(len(new_cache)):
                 for kv_idx in range(len(new_cache[layer_idx])):
                     self.assertTrue(