Update modeling_chatglm.py

modeling_chatglm.py

@@ -42,6 +42,7 @@ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
 _CONFIG_FOR_DOC = "ChatGLMConfig"
 
 is_transformers_4_42_or_higher = int(transformers.__version__.split(".")[1]) >= 42
+is_transformers_4_44_or_higher = int(transformers.__version__.split(".")[1]) >= 44
 
 
 def default_init(cls, *args, **kwargs):
@@ -812,8 +813,13 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             is_encoder_decoder: bool = False,
             standardize_cache_format: bool = False,
     ) -> Dict[str, Any]:
+
+        if is_transformers_4_44_or_higher:
+            model_kwargs["past_key_values"] = self._extract_past_from_model_output(
+                outputs
+            )[1]
         # update past_key_values
-        if is_transformers_4_42_or_higher:
+        elif is_transformers_4_42_or_higher:
             model_kwargs["past_key_values"] = self._extract_past_from_model_output(
                 outputs, standardize_cache_format=standardize_cache_format
             )[1]