Predict the lbl for particular pet :param batch: :param pet: :return:
(self, batch)
| 129 | return loss |
| 130 | |
| 131 | def predict(self, batch): |
| 132 | """ |
| 133 | Predict the lbl for particular pet |
| 134 | :param batch: |
| 135 | :param pet: |
| 136 | :return: |
| 137 | """ |
| 138 | if self.config.model_modifier == "intrinsic": |
| 139 | intrinsic_plugin_on_step(self) |
| 140 | |
| 141 | input_ids, choices_ids, labels = batch["input_ids"], batch["answer_choices_ids"], batch["labels"] |
| 142 | |
| 143 | if not self.config.split_option_at_inference: |
| 144 | bs, num_choices = choices_ids.size()[:2] |
| 145 | flat_choices_ids = choices_ids.flatten(0, 1) |
| 146 | attention_mask = (input_ids != self.tokenizer.pad_token_id).float() # [bs, max_seq_len] |
| 147 | encoder_hidden_states = self.model.encoder(input_ids=input_ids, attention_mask=attention_mask)[0] |
| 148 | encoder_hidden_states = encoder_hidden_states.unsqueeze(dim=1).repeat(1, num_choices, 1, 1).flatten(0, 1) |
| 149 | attention_mask = attention_mask.unsqueeze(dim=1).repeat(1, num_choices, 1).flatten(0, 1) |
| 150 | decoder_input_ids = torch.cat([torch.zeros_like(flat_choices_ids[:, :1]), flat_choices_ids[:, :-1]], dim=1) |
| 151 | decoder_attention_mask = (decoder_input_ids == decoder_input_ids).float() |
| 152 | lm_target = flat_choices_ids - 100 * (flat_choices_ids == self.tokenizer.pad_token_id).long() |
| 153 | |
| 154 | model_output = self.model( |
| 155 | attention_mask=attention_mask, |
| 156 | encoder_outputs=[encoder_hidden_states], |
| 157 | decoder_input_ids=decoder_input_ids, |
| 158 | decoder_attention_mask=decoder_attention_mask, |
| 159 | ) |
| 160 | choices_scores = ( |
| 161 | F.cross_entropy(model_output.logits.flatten(0, 1), lm_target.flatten(0, 1), reduction="none") |
| 162 | .view(bs, num_choices, -1) |
| 163 | .sum(dim=-1) |
| 164 | ) |
| 165 | # Length normalization |
| 166 | if self.config.length_norm > 0: |
| 167 | choices_scores = choices_scores / torch.pow( |
| 168 | (choices_ids != self.tokenizer.pad_token_id).sum(dim=-1), self.config.length_norm |
| 169 | ) |
| 170 | pred_score, prediction = choices_scores.min(dim=1) |
| 171 | |
| 172 | else: |
| 173 | bs, num_choices = choices_ids.size()[:2] |
| 174 | midpoint = num_choices // 2 |
| 175 | # |
| 176 | first_half_choice_ids = choices_ids[:, :midpoint, :] |
| 177 | second_half_choice_ids = choices_ids[:, midpoint:, :] |
| 178 | # |
| 179 | all_choice_scores = [] |
| 180 | |
| 181 | for half_choice_ids in [first_half_choice_ids, second_half_choice_ids]: |
| 182 | half_num_choices = half_choice_ids.shape[1] |
| 183 | |
| 184 | flat_choices_ids = half_choice_ids.flatten(0, 1) # [bs*num_choices, choice_len] |
| 185 | |
| 186 | attention_mask = (input_ids != self.tokenizer.pad_token_id).float() # [bs, max_seq_len] |
| 187 | encoder_hidden_states = self.model.encoder(input_ids=input_ids, attention_mask=attention_mask)[0] |
| 188 | encoder_hidden_states = ( |
no outgoing calls