examples/research_projects/onnx/summarization/bart_onnx/generation_onnx.py/BeamSearchScorerTS/init
class BeamSearchScorerTS:    def init(
        self,
        batch_size: int,
        max_length: int,
        num_beams: int,
        device: torch.device,
        length_penalty: float = 1.0,
        do_early_stopping: bool = False,
        num_beam_hyps_to_keep: int = 1,
        num_beam_groups: int = 1,
    ):
        self.max_length = max_length
        self.num_beams = num_beams
        self.batch_size = batch_size
        self.length_penalty = length_penalty
        self.do_early_stopping = do_early_stopping
        self.num_beam_hyps_to_keep = num_beam_hyps_to_keep
        self.num_beam_groups = num_beam_groups
        self.group_size = self.num_beams // self.num_beam_groups

        # NOTE: TorchScript does not support List of Modules
        #       Rewritten BeamHypotheses with tensors and list of tensors.
        self._done = torch.zeros(batch_size, dtype=torch.bool, device=device)
        self._beam_hyps_count = torch.zeros(batch_size, dtype=torch.long, device=device)
        self._beam_hyps_worst_scores = torch.zeros(batch_size, device=device) + 1e9
        self._beam_hyps = []
        self._beam_scores = []

        self._beam_hyps_max_length = max_length - 1  # ignoring bos_token

        if not isinstance(num_beams, int) or num_beams <= 1:
            raise ValueError(
                f"`num_beams` has to be an integer strictly greater than 1, but is {num_beams}. For `num_beams` == 1,"
                " one should make use of `greedy_search` instead."
            )

        if not isinstance(num_beam_groups, int) or (num_beam_groups > num_beams) or (num_beams % num_beam_groups != 0):
            raise ValueError(
                "`num_beam_groups` has to be an integer smaller or equal than `num_beams` and `num_beams` has to be"
                f" divisible by `num_beam_groups`, but is {num_beam_groups} with `num_beams` being {num_beams}."
            )