| 849 | |
| 850 | |
| 851 | class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel): |
| 852 | def __init__(self, config: ChatGLMConfig, empty_init=True, device=None): |
| 853 | super().__init__(config) |
| 854 | |
| 855 | self.max_sequence_length = config.max_length |
| 856 | self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device) |
| 857 | self.config = config |
| 858 | self.quantized = False |
| 859 | |
| 860 | if self.config.quantization_bit: |
| 861 | self.quantize(self.config.quantization_bit, empty_init=True) |
| 862 | |
| 863 | def _update_model_kwargs_for_generation( |
| 864 | self, |
| 865 | outputs: ModelOutput, |
| 866 | model_kwargs: Dict[str, Any], |
| 867 | is_encoder_decoder: bool = False, |
| 868 | standardize_cache_format: bool = False, |
| 869 | ) -> Dict[str, Any]: |
| 870 | # update past_key_values |
| 871 | model_kwargs["past_key_values"] = self._extract_past_from_model_output( |
| 872 | outputs, standardize_cache_format=standardize_cache_format |
| 873 | ) |
| 874 | |
| 875 | # update attention mask |
| 876 | if "attention_mask" in model_kwargs: |
| 877 | attention_mask = model_kwargs["attention_mask"] |
| 878 | model_kwargs["attention_mask"] = torch.cat( |
| 879 | [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 |
| 880 | ) |
| 881 | |
| 882 | # update position ids |
| 883 | if "position_ids" in model_kwargs: |
| 884 | position_ids = model_kwargs["position_ids"] |
| 885 | new_position_id = position_ids[..., -1:].clone() |
| 886 | new_position_id += 1 |
| 887 | model_kwargs["position_ids"] = torch.cat( |
| 888 | [position_ids, new_position_id], dim=-1 |
| 889 | ) |
| 890 | |
| 891 | model_kwargs["is_first_forward"] = False |
| 892 | return model_kwargs |
| 893 | |
| 894 | def prepare_inputs_for_generation( |
| 895 | self, |
| 896 | input_ids: torch.LongTensor, |
| 897 | past_key_values: Optional[torch.Tensor] = None, |
| 898 | attention_mask: Optional[torch.Tensor] = None, |
| 899 | position_ids: Optional[torch.Tensor] = None, |
| 900 | use_cache: Optional[bool] = None, |
| 901 | is_first_forward: bool = True, |
| 902 | **kwargs |
| 903 | ) -> dict: |
| 904 | # only last token for input_ids if past is not None |
| 905 | if position_ids is None: |
| 906 | position_ids = self.get_position_ids(input_ids, device=input_ids.device) |
| 907 | if not is_first_forward: |
| 908 | if past_key_values is not None: |
nothing calls this directly
no outgoing calls
no test coverage detected