Update modeling_chatglm.py

modeling_chatglm.py

@@ -47,6 +47,7 @@ CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
 ]
 
 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):
@@ -870,7 +871,11 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             is_encoder_decoder: bool = False,
             standardize_cache_format: bool = False,
     ) -> Dict[str, Any]:
-        if is_transformers_4_42_or_higher:
+        if is_transformers_4_44_or_higher:
+            model_kwargs["past_key_values"] = self._extract_past_from_model_output(
+                outputs
+            )[1]
+        elif is_transformers_4_42_or_higher:
             # update past_key_values
             model_kwargs["past_key_values"] = self._extract_past_from_model_output(
                 outputs, standardize_cache_format=standardize_cache_format