(self, language_model, take_softmax=True, length_penalty=0.0)
| 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): |