| 6 | |
| 7 | |
| 8 | class GLMForMultiTokenCloze(torch.nn.Module): |
| 9 | def __init__(self, language_model: GLMModel, take_softmax=True, length_penalty=0.0): |
| 10 | super(GLMForMultiTokenCloze, self).__init__() |
| 11 | self.model = language_model |
| 12 | self.take_softmax = take_softmax |
| 13 | self.length_penalty = length_penalty |
| 14 | |
| 15 | def state_dict(self, destination=None, prefix='', keep_vars=False): |
| 16 | # [h.remove() for h in self.hook_handles] |
| 17 | sd = self.model.state_dict(destination, prefix, keep_vars) |
| 18 | return sd |
| 19 | |
| 20 | def load_state_dict(self, state_dict, strict=True): |
| 21 | return self.model.load_state_dict(state_dict, strict=strict) |
| 22 | |
| 23 | def named_parameters(self, prefix: str = '', recurse: bool = True): |
| 24 | return self.model.named_parameters(prefix=prefix, recurse=recurse) |
| 25 | |
| 26 | def forward(self, input_ids, position_ids, attention_mask, target_ids=None, logit_mask=None, prompt_pos=None): |
| 27 | if target_ids == None: |
| 28 | return self.model(input_ids, position_ids, attention_mask) |
| 29 | num_choices = None |
| 30 | if len(input_ids.shape) == 3: |
| 31 | batch_size, num_choices = input_ids.shape[:2] |
| 32 | input_ids = input_ids.reshape(-1, input_ids.size(-1)) |
| 33 | attention_mask = attention_mask.reshape(-1, *attention_mask.size()[2:]) |
| 34 | position_ids = position_ids.reshape(-1, *position_ids.size()[2:]) |
| 35 | target_ids = target_ids.reshape(-1, target_ids.size(-1)) |
| 36 | logit_mask = logit_mask.reshape(-1, logit_mask.size(-1)) |
| 37 | if prompt_pos is not None: |
| 38 | prompt_pos = prompt_pos.reshape(-1, prompt_pos.size(-1)) |
| 39 | outputs, *mems = self.model(input_ids, position_ids, attention_mask, prompt_pos=prompt_pos) |
| 40 | if self.take_softmax: |
| 41 | outputs = torch.nn.functional.log_softmax(outputs, dim=-1) |
| 42 | # select the target logits |
| 43 | batch_ids = torch.arange(target_ids.size(0), dtype=torch.long, device=target_ids.device) |
| 44 | batch_ids = batch_ids.unsqueeze(1).expand_as(target_ids) |
| 45 | seq_ids = torch.arange(target_ids.size(-1), dtype=torch.long, device=target_ids.device) |
| 46 | seq_ids = seq_ids.unsqueeze(0).expand_as(target_ids) |
| 47 | logits = outputs[batch_ids, seq_ids, target_ids] |
| 48 | logits = (logits * logit_mask).sum(dim=1) |
| 49 | if self.length_penalty > 0.0: |
| 50 | logits = logits / logit_mask.sum(dim=1) ** self.length_penalty |
| 51 | if num_choices is not None: |
| 52 | logits = logits.view(-1, num_choices) |
| 53 | return (logits, *mems) |
| 54 | |
| 55 | |
| 56 | class GLMForMultiTokenClozeFast(torch.nn.Module): |