(self, input_ids, position_ids, attention_mask, target_ids=None, logit_mask=None, prompt_pos=None)
| 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): |
nothing calls this directly
no outgoing calls
no test coverage detected