compatibility with new transformers (#60)

- compatibility with new transformers (9e21dac48837929ca4df28e3dcb6ae04c184573d)
- Update modeling_chatglm.py (c59cdd3bafd43d5f9b3e82c88c088eccb3925e02)


Co-authored-by: Ekaterina Aidova <[email protected]>

modeling_chatglm.py

@@ -14,6 +14,7 @@ from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
 from torch.nn.utils import skip_init
 from typing import Optional, Tuple, Union, List, Callable, Dict, Any
 from copy import deepcopy
+import transformers
 
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
@@ -45,6 +46,9 @@ CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
     # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
 ]
 
+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):
     return cls(*args, **kwargs)
@@ -872,9 +876,19 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             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 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
+            )[1]
+        else:
+            model_kwargs["past_key_values"] = self._extract_past_from_model_output(
+                outputs, standardize_cache_format=standardize_cache_format
+            )
 
         # update attention mask
         if "attention_mask" in model_kwargs: