| 54 | |
| 55 | |
| 56 | class GLMForMultiTokenClozeFast(torch.nn.Module): |
| 57 | def __init__(self, language_model, take_softmax=True, length_penalty=0.0): |
| 58 | super(GLMForMultiTokenClozeFast, self).__init__() |
| 59 | self.model = language_model |
| 60 | self.take_softmax = take_softmax |
| 61 | self.length_penalty = length_penalty |
| 62 | |
| 63 | def forward(self, input_ids, position_ids, attention_mask, |
| 64 | dec_input_ids, dec_position_ids, dec_attention_mask, dec_target_ids, dec_logit_mask): |
| 65 | # encoder |
| 66 | outputs, *mems = self.model(input_ids, position_ids, attention_mask, return_memory=True, detach_memory=False) |
| 67 | batch_size, num_choices, max_dec_len = dec_input_ids.size() |
| 68 | max_enc_len = input_ids.size(-1) |
| 69 | |
| 70 | enc_mems = [] |
| 71 | for hidden in mems: |
| 72 | hidden = hidden.unsqueeze(1).expand(-1, num_choices, -1, -1).reshape(batch_size * num_choices, |
| 73 | *hidden.size()[1:]) |
| 74 | enc_mems.append(hidden) |
| 75 | |
| 76 | def build_dec_mask_matrix(seq_length, sep, memory_length=0): |
| 77 | m = enc_mems[0].new_ones((1, seq_length, seq_length)) |
| 78 | m = torch.tril(m) |
| 79 | |
| 80 | # sep = dec_attention_mask |
| 81 | ids = torch.arange(memory_length, device=sep.device, dtype=sep.dtype).view(1, -1) |
| 82 | mask = ids < sep.view(-1, 1) # batch * mem |
| 83 | mask = mask.unsqueeze(1).float().expand(-1, seq_length, -1) |
| 84 | |
| 85 | m = m.expand(batch_size * num_choices, -1, -1) |
| 86 | m = torch.cat((mask, m), dim=2) |
| 87 | m = m.unsqueeze(1) |
| 88 | return m |
| 89 | |
| 90 | dec_input_ids = dec_input_ids.reshape(-1, max_dec_len) |
| 91 | dec_position_ids = dec_position_ids.reshape(-1, *dec_position_ids.size()[2:]) |
| 92 | # dec_attention_mask = dec_attention_mask.reshape(-1, *dec_attention_mask.size()[2:]).unsqueeze(1) |
| 93 | dec_attention_mask = build_dec_mask_matrix(max_dec_len, dec_attention_mask.reshape(-1), max_enc_len) |
| 94 | dec_target_ids = dec_target_ids.reshape(-1, dec_target_ids.size(-1)) |
| 95 | dec_logit_mask = dec_logit_mask.reshape(-1, dec_logit_mask.size(-1)) |
| 96 | |
| 97 | outputs, *mems = self.model(dec_input_ids, dec_position_ids, dec_attention_mask, *enc_mems) |
| 98 | if self.take_softmax: |
| 99 | outputs = torch.nn.functional.log_softmax(outputs, dim=-1) |
| 100 | |
| 101 | batch_ids = torch.arange(dec_target_ids.size(0), dtype=torch.long, device=dec_target_ids.device) |
| 102 | batch_ids = batch_ids.unsqueeze(1).expand_as(dec_target_ids) |
| 103 | seq_ids = torch.arange(dec_target_ids.size(-1), dtype=torch.long, device=dec_target_ids.device) |
| 104 | seq_ids = seq_ids.unsqueeze(0).expand_as(dec_target_ids) |
| 105 | logits = outputs[batch_ids, seq_ids, dec_target_ids] |
| 106 | logits = (logits * dec_logit_mask).sum(dim=1) |
| 107 | if self.length_penalty > 0.0: |
| 108 | logits = logits / dec_logit_mask.sum(dim=1) ** self.length_penalty |
| 109 | if num_choices is not None: |
| 110 | logits = logits.view(-1, num_choices) |
| 111 | return (logits, *mems) |
| 112 | |
| 113 | |