| 27 | from transformers import AutoConfig |
| 28 | |
| 29 | class ChatModel(nn.Module, GenerationMixin): |
| 30 | def __init__(self, args, model=None): |
| 31 | super().__init__() |
| 32 | self.config = AutoConfig.from_pretrained('THUDM/chatglm2-6b', trust_remote_code=True) |
| 33 | self.generation_config = GenerationConfig.from_model_config(self.config) |
| 34 | if model is None: |
| 35 | self.model, self.args = AutoModel.from_pretrained("chatglm2-6b", args) |
| 36 | else: |
| 37 | self.model, self.args = model, args |
| 38 | self.device = self.model.parameters().__next__().device |
| 39 | self.main_input_name = 'input_ids' |
| 40 | |
| 41 | @classmethod |
| 42 | def from_pretrained(cls, name, args=None, base_cls=None, *, home_path=None, url=None, prefix='', **kwargs): |
| 43 | if base_cls is None: |
| 44 | model, args = AutoModel.from_pretrained(name, args, home_path=home_path, url=url, prefix=prefix, **kwargs) |
| 45 | else: |
| 46 | model, args = base_cls.from_pretrained(name, args, home_path=home_path, url=url, prefix=prefix, **kwargs) |
| 47 | return cls(args, model), args |
| 48 | |
| 49 | def can_generate(self): |
| 50 | return True |
| 51 | |
| 52 | def _update_model_kwargs_for_generation( |
| 53 | self, |
| 54 | outputs: ModelOutput, |
| 55 | model_kwargs: Dict[str, Any], |
| 56 | is_encoder_decoder: bool = False, |
| 57 | standardize_cache_format: bool = False, |
| 58 | ) -> Dict[str, Any]: |
| 59 | # update past_key_values |
| 60 | model_kwargs["past_key_values"] = self._extract_past_from_model_output( |
| 61 | outputs, standardize_cache_format=standardize_cache_format |
| 62 | ) |
| 63 | |
| 64 | # update attention mask |
| 65 | if "attention_mask" in model_kwargs: |
| 66 | attention_mask = model_kwargs["attention_mask"] |
| 67 | model_kwargs["attention_mask"] = torch.cat( |
| 68 | [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 |
| 69 | ) |
| 70 | |
| 71 | # update position ids |
| 72 | if "position_ids" in model_kwargs: |
| 73 | position_ids = model_kwargs["position_ids"] |
| 74 | new_position_id = position_ids[..., -1:].clone() |
| 75 | new_position_id += 1 |
| 76 | model_kwargs["position_ids"] = torch.cat( |
| 77 | [position_ids, new_position_id], dim=-1 |
| 78 | ) |
| 79 | |
| 80 | model_kwargs["is_first_forward"] = False |
| 81 | return model_kwargs |
| 82 | |
| 83 | def prepare_inputs_for_generation( |
| 84 | self, |
| 85 | input_ids: torch.LongTensor, |
| 86 | past: Optional[torch.Tensor] = None, |