examples/research_projects/onnx/summarization/bart_onnx/generation_onnx.py/BeamSearchScorerTS/finalize
class BeamSearchScorerTS:    def finalize(
        self,
        input_ids: torch.Tensor,
        final_beam_scores: torch.Tensor,
        final_beam_tokens: torch.Tensor,
        final_beam_indices: torch.Tensor,
        pad_token_id: int,
        eos_token_id: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        batch_size = len(self._beam_hyps_count)

        # finalize all open beam hypotheses and add to generated hypotheses
        for batch_idx in range(batch_size):
            if self._done[batch_idx]:
                continue

            # all open beam hypotheses are added to the beam hypothesis
            # beam hypothesis class automatically keeps the best beams
            for beam_id in range(self.num_beams):
                batch_beam_idx = batch_idx * self.num_beams + beam_id
                final_score = final_beam_scores[batch_beam_idx].item()
                final_tokens = input_ids[batch_beam_idx]
                self.hypo_add(final_tokens, final_score, batch_idx)

        # select the best hypotheses
        # NOTE: torch.Tensor.new_zeros() is not scriptable
        sent_lengths = torch.zeros(batch_size * self.num_beam_hyps_to_keep, dtype=torch.long)
        best = []
        best_scores = torch.zeros(
            batch_size * self.num_beam_hyps_to_keep, device=input_ids.device, dtype=torch.float32
        )
        # retrieve best hypotheses
        for i in range(batch_size):
            # NOTE: lambda is not scriptable
            batch_hypo_start = torch.sum(self._beam_hyps_count[:i]) if i > 0 else torch.tensor(0, dtype=torch.long)
            batch_hypo_end = torch.sum(self._beam_hyps_count[: i + 1])
            beam_scores = torch.cat(self._beam_scores)[batch_hypo_start:batch_hypo_end]
            sorted_next_scores, sorted_indices = torch.topk(beam_scores, len(beam_scores), largest=True)
            for j in range(self.num_beam_hyps_to_keep):
                best_score = beam_scores[sorted_indices[j]]
                best_hyp = self._beam_hyps[batch_hypo_start + sorted_indices[j]]
                sent_lengths[self.num_beam_hyps_to_keep * i + j] = len(best_hyp)
                # append to lists
                best.append(best_hyp)
                best_scores[i * self.num_beam_hyps_to_keep + j] = best_score

        # prepare for adding eos
        sent_max_len = min(sent_lengths.max() + 1, self.max_length)
        decoded = torch.zeros(batch_size * self.num_beam_hyps_to_keep, sent_max_len, dtype=torch.long)
        # shorter batches are padded if needed
        if sent_lengths.min() != sent_lengths.max():
            assert pad_token_id is not None, "`pad_token_id` has to be defined"
            decoded.fill_(pad_token_id)

        # fill with hypotheses and eos_token_id if the latter fits in
        for i, hypo in enumerate(best):
            decoded[i, : sent_lengths[i]] = hypo
            if sent_lengths[i] < self.max_length:
                decoded[i, sent_lengths[i]] = eos_token_id

        return decoded, best_scores