def _get_best_finished_and_nonfinished_candidates(self, expanded_beam_scores, expanded_beam_tokens):
        """
        Get top-k next tokens for each candidate.
        Also updates aggregated scores according to the scores for chosen tokens.

        :param expanded_beam_scores: Aggregated scores for each response in the extended beam
        :param expanded_beam_tokens: Last tokens in each response in the extended beam
        :return: best_non_finished_candidates, best_finished_candidates
            best_non_finished_candidates are used for updating beam for the next step
            best_finished_candidates are returned to dump on each step and rerank afterwards.
        """
        # This mask contains true if the corresponding candidate is finished with <EOS> and false otherwise
        finished_candidates_mask = (expanded_beam_tokens == self._service_tokens_ids.eos_token_id)
        # We select the best candidates among those who are not finished
        best_non_finished_candidates, self._cur_candidates_scores = \
            self._get_k_max_elements_indices_and_scores(expanded_beam_scores, self._beam_size,
                                                        ~finished_candidates_mask)
        # And also return finished candidates that are good enough to fit in the new beam, but do not go there
        high_quality_candidates_mask = (expanded_beam_scores > self._cur_candidates_scores[self._beam_size - 1] - 1e-6)
        # We need [0] in the end of this line because np.nonzero returns tuple, but we only need indices
        best_finished_candidates = np.nonzero(finished_candidates_mask & high_quality_candidates_mask)[0]
        return best_non_finished_candidates, best_finished_candidates

        Get top-k next tokens for each candidate.
        Also updates aggregated scores according to the scores for chosen tokens.

        :param expanded_beam_scores: Aggregated scores for each response in the extended beam
        :param expanded_beam_tokens: Last tokens in each response in the extended beam
        :return: best_non_finished_candidates, best_finished_candidates
            best_non_finished_candidates are used for updating beam for the next step
            best_finished_candidates are returned to dump on each step and rerank afterwards.
        
cakechat/dialog_model/inference/candidates/beamsearch.py