zR
commited on
Update modeling_chatglm.py
Browse files- modeling_chatglm.py +2 -3
modeling_chatglm.py
CHANGED
|
@@ -870,9 +870,8 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
|
|
| 870 |
standardize_cache_format: bool = False,
|
| 871 |
) -> Dict[str, Any]:
|
| 872 |
# update past_key_values
|
| 873 |
-
|
| 874 |
-
|
| 875 |
-
)
|
| 876 |
|
| 877 |
# update attention mask
|
| 878 |
if "attention_mask" in model_kwargs:
|
|
|
|
| 870 |
standardize_cache_format: bool = False,
|
| 871 |
) -> Dict[str, Any]:
|
| 872 |
# update past_key_values
|
| 873 |
+
cache_name, cache = self._extract_past_from_model_output(outputs)
|
| 874 |
+
model_kwargs[cache_name] = cache
|
|
|
|
| 875 |
|
| 876 |
# update attention mask
|
| 877 |
if "attention_mask" in model_kwargs:
|