Upload modeling_chatglm.py

modeling_chatglm.py

@@ -827,7 +827,9 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             init_method = default_init
         init_kwargs = {}
         if device is not None:
-            init_kwargs["device"] = device
+            init_kwargs["device"] = torch.device(device)
+        if isinstance(config.torch_dtype, str):
+            config.torch_dtype = getattr(torch, config.torch_dtype)
         self.embedding = init_method(Embedding, config, **init_kwargs)
         self.num_layers = config.num_layers
         self.multi_query_group_num = config.multi_query_group_num
@@ -923,10 +925,17 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             outputs: ModelOutput,
             model_kwargs: Dict[str, Any],
             is_encoder_decoder: bool = False,
+            num_new_tokens: int = 1,
     ) -> Dict[str, Any]:
         # update past_key_values
-        cache_name, cache = self._extract_past_from_model_output(outputs)
-        model_kwargs[cache_name] = cache
+        for possible_cache_name in ["past_key_values", "mems", "past_buckets_states", "cache_params"]:
+            if hasattr(outputs, possible_cache_name):
+                if possible_cache_name in ("past_buckets_states", "mems"):
+                    cache_name = "past_key_values"
+                else:
+                    cache_name = possible_cache_name
+                model_kwargs[cache_name] = getattr(outputs, possible_cache_name)
+                break
 
         # update attention mask
         if "attention_mask" in model_kwargs:
@@ -945,6 +954,10 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             )
 
         model_kwargs["is_first_forward"] = False
+
+        if model_kwargs.get("use_cache", True) and "cache_position" in model_kwargs:
+            model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens
+
         return model_kwargs
 
     def prepare_inputs_for_generation(