Update custom_generate/generate.py

custom_generate/generate.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn as nn
 
 from transformers import GenerationConfig, LogitsProcessorList, StoppingCriteriaList
-from transformers.cache_utils import Cache, EncoderDecoderCache
+from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache
 from transformers.configuration_utils import PretrainedConfig
 from transformers.generation.utils import (
     ALL_CACHE_NAMES,
@@ -249,13 +249,17 @@ def _contrastive_search(
                     f"{model.__class__.__name__} does not support caching and therefore **can't** be used "
                     "for contrastive search."
                 )
-            elif (
-                not isinstance(past_key_values[0], (tuple, torch.Tensor))
-                or past_key_values[0][0].shape[0] != batch_size
+            # Only those caches have the necesary methods
+            elif not (
+                isinstance(past_key_values, DynamicCache)
+                or (
+                    isinstance(past_key_values, EncoderDecoderCache)
+                    and isinstance(past_key_values.self_attention_cache, DynamicCache)
+                )
             ):
                 raise ValueError(
-                    f"{model.__class__.__name__} does not have a standard cache format and therefore **can't** be "
-                    "used for contrastive search without further modifications."
+                    f"Unsupported cache type: {type(outputs['past_key_values'])}. Contrastive search requires "
+                    "dynamic cache, so set `cache_implementation='dynamic'` in the generation config."
                 )
 
         # contrastive_search main logic start: