def __init__(self, nn_model, beam_size, repetition_penalization_coefficient):
        """
        :param nn_model: NN model to use for predicting
        :param beam_size: Size of beam
        """
        self._nn_model = nn_model
        self._beam_size = min(beam_size, self._nn_model.vocab_size)
        self._num_finished_candidates_to_keep = self._beam_size ** 2
        self._log_repetition_penalization_coefficient = np.log(repetition_penalization_coefficient)
        self._service_tokens_ids = ServiceTokensIDs(nn_model.token_to_index)

        :param nn_model: NN model to use for predicting
        :param beam_size: Size of beam
        
cakechat/dialog_model/inference/candidates/beamsearch.py