| 84 | from .chatglm_model import ChatGLMFinalMixin |
| 85 | |
| 86 | class ChatGLM3Model(BaseModel): |
| 87 | def __init__(self, args, transformer=None, **kwargs): |
| 88 | super(ChatGLM3Model, self).__init__(args, transformer=transformer, activation_func=F.silu, layernorm=RMSNorm, **kwargs) |
| 89 | del self.transformer.position_embeddings |
| 90 | self.add_mixin("chatglm-final", ChatGLMFinalMixin(args.vocab_size, args.hidden_size)) |
| 91 | self.add_mixin("attn", ChatGLM3AttnMixin(args.hidden_size, args.num_attention_heads, args.max_sequence_length, args.base_scale)) |
| 92 | self.add_mixin("mlp", SwiGLUMixin(args.num_layers, args.hidden_size, args.inner_hidden_size, bias=args.use_bias)) |
| 93 | |
| 94 | def position_embedding_forward(self, position_ids, output_cross_layer, **kw_args): |
| 95 | return None |
| 96 | |
| 97 | def get_masks(self, input_ids, past_key_values, padding_mask=None): |
| 98 | batch_size, seq_length = input_ids.shape |
| 99 | full_attention_mask = torch.ones(batch_size, seq_length, seq_length, dtype=next(self.parameters()).dtype, device=input_ids.device) |
| 100 | full_attention_mask.tril_() |
| 101 | past_length = 0 |
| 102 | if past_key_values: |
| 103 | past_length = past_key_values[0][0].shape[2] |
| 104 | if past_length: |
| 105 | full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length, dtype=next(self.parameters()).dtype, |
| 106 | device=input_ids.device), full_attention_mask), dim=-1) |
| 107 | if padding_mask is not None: |
| 108 | full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1) |
| 109 | if not past_length and padding_mask is not None: |
| 110 | full_attention_mask -= padding_mask.unsqueeze(-1) - 1 |
| 111 | full_attention_mask = (full_attention_mask < 0.5).bool() |
| 112 | full_attention_mask.unsqueeze_(1) |
| 113 | return full_attention_mask |
| 114 | |
| 115 | def get_position_ids(self, input_ids): |
| 116 | batch_size, seq_length = input_ids.shape |
| 117 | position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device).unsqueeze(0).repeat(batch_size, 1) |
| 118 | return position_ids |
| 119 | |
| 120 | def forward(self, input_ids, position_ids=None, attention_mask=None, past_key_values=None, **kwargs): |
| 121 | if position_ids is None: |
| 122 | position_ids = self.get_position_ids(input_ids) |
| 123 | if attention_mask is not None and attention_mask.ndim == 4: |
| 124 | pass |
| 125 | elif past_key_values is not None and input_ids.size(0) == 1: |
| 126 | attention_mask = torch.tensor([[1]], dtype=torch.long, device=input_ids.device) |
| 127 | else: |
| 128 | attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask) |
| 129 | if attention_mask is not None and attention_mask.dtype is torch.bool: |
| 130 | attention_mask = ~attention_mask |
| 131 | attention_mask = attention_mask.to(next(self.parameters()).dtype) |
| 132 | if past_key_values is not None: |
| 133 | input_ids = input_ids[:, -1:] |
| 134 | position_ids = position_ids[..., -1:] |
| 135 | if input_ids.size(0) != 1: |
| 136 | attention_mask = attention_mask[:, :, -1:] |
| 137 | return super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, **kwargs) |
| 138 | |
| 139 | @classmethod |
| 140 | def add_model_specific_args(cls, parser): |
| 141 | group = parser.add_argument_group('ChatGLM3', 'ChatGLM3 Configurations') |
| 142 | group.add_argument('--base-scale', type=float, default=1.) |
| 143 | return super().add_model_specific_args(parser) |