| 642 | |
| 643 | |
| 644 | class ChatGLMModel(ChatGLMPreTrainedModel): |
| 645 | def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): |
| 646 | super().__init__(config) |
| 647 | if empty_init: |
| 648 | init_method = skip_init |
| 649 | else: |
| 650 | init_method = default_init |
| 651 | init_kwargs = {} |
| 652 | if device is not None: |
| 653 | init_kwargs["device"] = device |
| 654 | self.embedding = init_method(Embedding, config, **init_kwargs) |
| 655 | |
| 656 | # Rotary positional embeddings |
| 657 | self.seq_length = config.seq_length |
| 658 | rotary_dim = ( |
| 659 | config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels |
| 660 | ) |
| 661 | |
| 662 | self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio, original_impl=config.original_rope, |
| 663 | device=device, dtype=config.torch_dtype) |
| 664 | self.encoder = init_method(GLMTransformer, config, **init_kwargs) |
| 665 | self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False, |
| 666 | dtype=config.torch_dtype, **init_kwargs) |
| 667 | |
| 668 | def get_input_embeddings(self): |
| 669 | return self.embedding.word_embeddings |
| 670 | |
| 671 | def forward( |
| 672 | self, |
| 673 | input_ids, |
| 674 | position_ids: Optional[torch.Tensor] = None, |
| 675 | attention_mask: Optional[torch.BoolTensor] = None, |
| 676 | full_attention_mask: Optional[torch.BoolTensor] = None, |
| 677 | past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None, |
| 678 | inputs_embeds: Optional[torch.Tensor] = None, |
| 679 | use_cache: Optional[bool] = None, |
| 680 | output_hidden_states: Optional[bool] = None, |
| 681 | return_dict: Optional[bool] = None, |
| 682 | ): |
| 683 | output_hidden_states = ( |
| 684 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 685 | ) |
| 686 | use_cache = use_cache if use_cache is not None else self.config.use_cache |
| 687 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 688 | |
| 689 | batch_size, seq_length = input_ids.shape |
| 690 | |
| 691 | if inputs_embeds is None: |
| 692 | inputs_embeds = self.embedding(input_ids) |
| 693 | |
| 694 | # if full_attention_mask is None: |
| 695 | # if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1): |
| 696 | # full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) |
| 697 | |
| 698 | # Rotary positional embeddings |
| 699 | rotary_pos_emb = self.rotary_pos_emb(self.seq_length) |
| 700 | if position_ids is not None: |
| 701 | rotary_pos_emb = rotary_pos_emb[position_ids] |