| 999 | |
| 1000 | |
| 1001 | class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel): |
| 1002 | def __init__(self, config): |
| 1003 | super().__init__(config) |
| 1004 | |
| 1005 | # self.hidden_size = config.hidden_size |
| 1006 | # self.params_dtype = torch.half |
| 1007 | # self.vocab_size = config.vocab_size |
| 1008 | self.max_sequence_length = config.max_sequence_length |
| 1009 | |
| 1010 | self.position_encoding_2d = config.position_encoding_2d |
| 1011 | |
| 1012 | self.transformer = ChatGLMModel(config) |
| 1013 | |
| 1014 | self.lm_head = skip_init( |
| 1015 | nn.Linear, |
| 1016 | config.hidden_size, |
| 1017 | config.vocab_size, |
| 1018 | bias=False, |
| 1019 | dtype=torch.half |
| 1020 | ) |
| 1021 | |
| 1022 | def get_output_embeddings(self): |
| 1023 | return self.lm_head |
| 1024 | |
| 1025 | def set_output_embeddings(self, new_embeddings): |
| 1026 | self.lm_head = new_embeddings |
| 1027 | |
| 1028 | def get_masks_and_position_ids(self, seq, mask_position, context_length, device, gmask=False): |
| 1029 | attention_mask = torch.ones((1, context_length, context_length), device=device) |
| 1030 | attention_mask.tril_() |
| 1031 | attention_mask[..., :mask_position - 1] = 1 |
| 1032 | attention_mask.unsqueeze_(1) |
| 1033 | attention_mask = (attention_mask < 0.5).bool() |
| 1034 | |
| 1035 | if self.position_encoding_2d: |
| 1036 | seq_length = seq.index(150004) |
| 1037 | position_ids = torch.arange(context_length, dtype=torch.long, device=device) |
| 1038 | if not gmask: |
| 1039 | position_ids[seq_length:] = mask_position |
| 1040 | block_position_ids = torch.cat(( |
| 1041 | torch.zeros(seq_length, dtype=torch.long, device=device), |
| 1042 | torch.arange(context_length - seq_length, dtype=torch.long, device=device) + 1 |
| 1043 | )) |
| 1044 | position_ids = torch.stack((position_ids, block_position_ids), dim=0) |
| 1045 | else: |
| 1046 | position_ids = torch.arange(context_length, dtype=torch.long, device=device) |
| 1047 | if not gmask: |
| 1048 | position_ids[context_length - 1:] = mask_position |
| 1049 | |
| 1050 | position_ids = position_ids.unsqueeze(0) |
| 1051 | |
| 1052 | return attention_mask, position_ids |
| 1053 | |
| 1054 | def prepare_inputs_for_generation( |
| 1055 | self, |
| 1056 | input_ids: torch.LongTensor, |
| 1057 | past: Optional[torch.Tensor] = None, |
| 1058 | past_key_values: Optional[torch.Tensor] = None, |
nothing calls this directly
no outgoing calls
no test coverage detected