| 989 | |
| 990 | |
| 991 | class ChatGLMModel(ChatGLMPreTrainedModel): |
| 992 | def __init__(self, config: ChatGLMConfig, device=None, empty_init=True): |
| 993 | super().__init__(config) |
| 994 | if empty_init: |
| 995 | init_method = skip_init |
| 996 | else: |
| 997 | init_method = default_init |
| 998 | init_kwargs = {} |
| 999 | if device is not None: |
| 1000 | init_kwargs["device"] = device |
| 1001 | self.embedding = init_method(Embedding, config, **init_kwargs) |
| 1002 | self.num_layers = config.num_layers |
| 1003 | self.multi_query_group_num = config.multi_query_group_num |
| 1004 | self.kv_channels = config.kv_channels |
| 1005 | |
| 1006 | # Rotary positional embeddings |
| 1007 | self.seq_length = config.seq_length |
| 1008 | rotary_dim = ( |
| 1009 | config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels |
| 1010 | ) |
| 1011 | |
| 1012 | self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device, |
| 1013 | dtype=config.torch_dtype) |
| 1014 | self.encoder = init_method(GLMTransformer, config, **init_kwargs) |
| 1015 | self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False, |
| 1016 | dtype=config.torch_dtype, **init_kwargs) |
| 1017 | self.pre_seq_len = config.pre_seq_len |
| 1018 | self.prefix_projection = config.prefix_projection |
| 1019 | if self.pre_seq_len is not None: |
| 1020 | for param in self.parameters(): |
| 1021 | param.requires_grad = False |
| 1022 | self.prefix_tokens = torch.arange(self.pre_seq_len).long() |
| 1023 | self.prefix_encoder = PrefixEncoder(config) |
| 1024 | self.dropout = torch.nn.Dropout(0.1) |
| 1025 | |
| 1026 | def get_input_embeddings(self): |
| 1027 | return self.embedding.word_embeddings |
| 1028 | |
| 1029 | def get_prompt(self, batch_size, device, dtype=torch.half): |
| 1030 | prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device) |
| 1031 | past_key_values = self.prefix_encoder(prefix_tokens).type(dtype) |
| 1032 | past_key_values = past_key_values.view( |
| 1033 | batch_size, |
| 1034 | self.pre_seq_len, |
| 1035 | self.num_layers * 2, |
| 1036 | self.multi_query_group_num, |
| 1037 | self.kv_channels |
| 1038 | ) |
| 1039 | # seq_len, b, nh, hidden_size |
| 1040 | past_key_values = self.dropout(past_key_values) |
| 1041 | past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2) |
| 1042 | return past_key_values |
| 1043 | |
| 1044 | def forward( |
| 1045 | self, |
| 1046 | input_ids, |
| 1047 | position_ids: Optional[torch.Tensor] = None, |
| 1048 | attention_mask: Optional[torch.BoolTensor] = None, |