Update of modeling_cogvlm.py for Transformers newer version

modeling_cogvlm.py

@@ -1,9 +1,11 @@
 """largely copy from llama and adapt for cogvlm"""
 import warnings
+import packaging.version
 from typing import TYPE_CHECKING, Optional, Tuple, List, Union, Literal, Dict, Any
 
 import math
 import torch
+import transformers
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from torchvision import transforms
@@ -26,7 +28,12 @@ logger = get_logger(__name__)
 
 LANGUAGE_TOKEN_TYPE = 0
 VISION_TOKEN_TYPE = 1
-
+TRANSFORMERS_ABOVE_441 = (
+    True
+    if packaging.version.parse(transformers.__version__)
+    >= packaging.version.parse("4.42.0")
+    else False
+)
 
 # Copied from transformers.models.bart.modeling_bart._make_causal_mask
 def _make_causal_mask(
@@ -736,9 +743,13 @@ class CogVLMForCausalLM(CogVLMPreTrainedModel):
             standardize_cache_format: bool = False,
     ) -> Dict[str, Any]:
         # update past_key_values
-        model_kwargs["past_key_values"] = self._extract_past_from_model_output(
-            outputs, standardize_cache_format=standardize_cache_format
-        )
+        if TRANSFORMERS_ABOVE_441:
+            cache_name, cache = self._extract_past_from_model_output(outputs)
+            model_kwargs[cache_name] = cache
+        else:
+            model_kwargs["past_key_values"] = self._extract_past_from_model_output(
+                outputs, standardize_cache_format=standardize_cache_format
+            )
         if getattr(outputs, "state", None) is not None:
             model_kwargs["state"] = outputs.state