Method evaluate

tasks/seq2seq/evaluate.py:267–390

Calculate correct over total answers and return prediction if the `output_predictions` is true.

(self, model, dataloader, example_dict, args)

            self.processors.append(processor)

    def evaluate(self, model, dataloader, example_dict, args):
        """Calculate correct over total answers and return prediction if the
        `output_predictions` is true."""
        model.eval()
        local_predictions = {}
        print_rank_0("Distributed store created")
        with torch.no_grad():
            # For all the batches in the dataset.
            for idx, data in enumerate(dataloader):
                tokens, attention_mask, position_ids = process_batch(data, args)
                batch_size = tokens.size(0)
                beam_scorer = BeamSearchScorer(
                    batch_size=batch_size,
                    max_length=args.out_seq_length,
                    num_beams=args.num_beams,
                    device=tokens.device,
                    length_penalty=args.length_penalty,
                    do_early_stopping=False,
                )
                beam_scores = torch.zeros((batch_size, args.num_beams), dtype=torch.float, device=tokens.device)
                beam_scores[:, 1:] = -1e9
                beam_scores = beam_scores.view((batch_size * args.num_beams,))
                # Run the model forward.
                counter = 0
                context_length = tokens.size(1)
                while counter < args.tgt_seq_length:
                    if counter == 0:
                        next_token_logits, *mems = model(tokens, position_ids, attention_mask, return_memory=True)
                        seq_length = next_token_logits.size(1)
                        next_token_logits = next_token_logits[:, -1]
                        next_token_logits = next_token_logits.unsqueeze(1).repeat(1, args.num_beams, 1).view(
                            batch_size * args.num_beams, -1)
                        mems = [mem.unsqueeze(1).repeat(1, args.num_beams, 1, 1).view(batch_size * args.num_beams,
                                                                                      seq_length, -1) for mem in mems]
                        position_ids = tokens.new_ones(batch_size, args.num_beams, 2, 1)
                        for i, text in enumerate(tokens.tolist()):
                            mask_pos = text.index(self.mask_token)
                            position_ids[i, :, 0] = mask_pos
                        position_ids = position_ids.reshape(batch_size * args.num_beams, 2, 1)
                        tokens = tokens.new_zeros(batch_size * args.num_beams, 0)
                    else:
                        if not args.no_block_position:
                            position_ids[:, 1] = counter + 1
                        last_token = tokens[:, -1:]
                        if self.mask_pad_token:
                            cur_attention_mask = attention_mask[:, :, -1:, :].unsqueeze(1).expand(-1, args.num_beams, -1,
                                                                                                  -1, -1).reshape(
                                batch_size * args.num_beams, 1, 1, context_length)
                            cur_attention_mask = torch.cat(
                                (cur_attention_mask, attention_mask.new_ones((batch_size * args.num_beams, 1, 1, counter))),
                                dim=-1)
                        else:
                            cur_attention_mask = tokens.new_zeros([batch_size * args.num_beams])
                        next_token_logits, *mems = model(last_token, position_ids, cur_attention_mask, *mems,
                                                         return_memory=True)
                        next_token_logits = next_token_logits[:, -1]
                    next_token_logits = top_k_logits(next_token_logits, top_k=args.top_k, top_p=args.top_p)
                    next_token_scores = F.log_softmax(next_token_logits, dim=-1)