Since `transformers` v4.56.0` the dictionary `ALL_STATIC_CACHE_IMPLEMENTATIONS` replaced `NEED_SETUP_CACHE_CLASSES_MAPPING`

modeling_decilm.py

@@ -19,15 +19,23 @@
 # limitations under the License.
 
 import math
+import importlib.metdata
 from typing import List, Optional, Tuple, Union
 
+from packaging.version import Version
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers import GenerationConfig
-from transformers.generation.utils import NEED_SETUP_CACHE_CLASSES_MAPPING, GenerationMixin, GenerateOutput
+from transformers.generation.utils import GenerationMixin, GenerateOutput
+
+
+if Version(importlib.metadata.version("transformers")) <= Version("4.56.0.dev0")
+  from transformers.generation.configuration_utils import NEED_SETUP_CACHE_CLASSES_MAPPING
+else:
+  from transformers.generation.configuration_utils import ALL_STATIC_CACHE_IMPLEMENTATIONS
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 from transformers.utils import (
@@ -810,7 +818,10 @@ class DeciLMPreTrainedModel(PreTrainedModel):
         # DeciLM-specific code
         generation_config, model_kwargs = super()._prepare_generation_config(generation_config, *args, **kwargs)
         generation_config.cache_implementation = "variable"
-        NEED_SETUP_CACHE_CLASSES_MAPPING["variable"] = VariableCache
+        if transformers_version <= Version("4.56.0.dev0")
+            NEED_SETUP_CACHE_CLASSES_MAPPING["variable"] = VariableCache
+        else:
+            ALL_STATIC_CACHE_IMPLEMENTATIONS["variable"] = VariableCache
         return generation_config, model_kwargs
 
 
@@ -1148,6 +1159,7 @@ class DeciLMForCausalLM(DeciLMPreTrainedModel, GenerationMixin):
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             cache_position: Optional[torch.LongTensor] = None,
+            **kwargs,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
         Args: