| 112 | |
| 113 | |
| 114 | class GLMForSingleTokenCloze(torch.nn.Module): |
| 115 | def __init__(self, language_model, take_softmax=False): |
| 116 | super().__init__() |
| 117 | self.model = language_model |
| 118 | self.take_softmax = take_softmax |
| 119 | |
| 120 | def state_dict(self, destination=None, prefix='', keep_vars=False): |
| 121 | # [h.remove() for h in self.hook_handles] |
| 122 | sd = self.model.state_dict(destination, prefix, keep_vars) |
| 123 | return sd |
| 124 | |
| 125 | def load_state_dict(self, state_dict, strict=True): |
| 126 | return self.model.load_state_dict(state_dict, strict=strict) |
| 127 | |
| 128 | def named_parameters(self, prefix: str = '', recurse: bool = True): |
| 129 | return self.model.named_parameters(prefix=prefix, recurse=recurse) |
| 130 | |
| 131 | def forward(self, input_ids, position_ids, attention_mask, target_ids=None, logit_mask=None, prompt_pos=None): |
| 132 | if target_ids is None: |
| 133 | return self.model(input_ids, position_ids, attention_mask) |
| 134 | assert len(input_ids.shape) == 2 |
| 135 | outputs, *mems = self.model(input_ids, position_ids, attention_mask, prompt_pos=prompt_pos) |
| 136 | batch_ids = torch.arange(outputs.size(0), dtype=attention_mask.dtype, device=attention_mask.device) |
| 137 | target_logits = outputs[batch_ids, attention_mask] |
| 138 | if self.take_softmax: |
| 139 | target_prob = torch.nn.functional.log_softmax(target_logits, dim=-1) |
| 140 | else: |
| 141 | target_prob = target_logits |
| 142 | batch_ids = batch_ids.unsqueeze(1).expand_as(target_ids) |
| 143 | output = target_prob[batch_ids, target_ids] |
| 144 | |
| 145 | return (output, target_logits, *mems) |
| 146 | |
| 147 | |
| 148 | class GLMForSequenceClassification(torch.nn.Module): |