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