def _update_next_candidates_and_hidden_states(self, token_idx, best_non_finished_candidates_indices,
                                                  expanded_beam_tokens):
        """
        Updates current state of candidates prediction process and fills in current candidates.

        :param token_idx: position of current token
        :param expanded_beam_tokens: np.array with shape (beam_size * beam_size,)
            Tokens candidates for the next step.
        :param best_non_finished_candidates_indices: np.array with shape (beam_size,)
            Contains indexes of best K candidates in current K^2 sized expanded beam to use in the next beam.
        """
        # Separate arrays for the updated hidden states
        next_hidden_states_batch = np.zeros_like(self._hidden_states_batch)
        # and the candidates
        next_step_candidates = np.full_like(self._cur_candidates, self._service_tokens_ids.pad_token_id, dtype=INTX)

        for i, candidate_idx in enumerate(best_non_finished_candidates_indices):
            # expanded_beam_tokens contains the last token for each of the beam_size^2 candidates in the expanded beam.
            # We need to get which original candidate this token in the expanded beam corresponds to.
            # (to fill in all the previous tokens from self._cur_candidates)
            # Because all the candidates in the expanded beam were filled sequentially, we just use this formula:
            original_candidate_idx = candidate_idx // self._beam_size

            # Construct the candidates for the next step using self._cur_candidates and the last token:

            # next_tokens is the last token for each new candidate here.
            next_token = expanded_beam_tokens[candidate_idx]

            # First, fill in all the preceding tokens for current candidate
            next_step_candidates[i, :token_idx] = self._cur_candidates[original_candidate_idx, :token_idx]
            # And put the last token of the current candidate on its position
            next_step_candidates[i, token_idx] = next_token
            # We also have to update the hidden states for the next step: we need to know which hidden states
            # to use to continue decoding each candidate and we get them from the corresponding positions
            next_hidden_states_batch[i] = self._hidden_states_batch[original_candidate_idx]

        self._hidden_states_batch = next_hidden_states_batch
        self._cur_candidates = next_step_candidates

        Updates current state of candidates prediction process and fills in current candidates.

        :param token_idx: position of current token
        :param expanded_beam_tokens: np.array with shape (beam_size * beam_size,)
            Tokens candidates for the next step.
        :param best_non_finished_candidates_indices: np.array with shape (beam_size,)
            Contains indexes of best K candidates in current K^2 sized expanded beam to use in the next beam.
        
cakechat/dialog_model/inference/candidates/beamsearch.py