def sample(self, probabilities, sample_idx, temperature=1.0):
        """
        Sample using individual priors for each sample_idx.
        Also updates individual priors for each sample in batch.
        We need individual priors to prevent the model from repeating the same tokens over and over again in one
        response.

        probabilities: Probabilities of each token. The distribution given by these probabilities
            is used by this function to sample the token.
        sample_idx : Integer between 0 and batch_size-1. We need it to figure out which token_log_prior line to use.
        temperature: Temperature for sampling. Temperature has to be a positive number.
        :return: Index of sampled token
        """
        # To make the repetition penalization invariant to the original temperature we have to adjust the coefficient:
        repetition_penalize_coefficient = np.exp(np.log(self._repetition_penalization_coefficient) / temperature)
        # Back-up the array to avoid side-effects (otherwise the function will change the probabilities passed as an
        # argument)
        probabilities = np.copy(probabilities)

        probabilities[self._banned_tokens_ids] = 0
        probabilities[self._used_tokens_ids[sample_idx]] /= repetition_penalize_coefficient

        probabilities /= np.sum(probabilities)
        token_id = np.random.choice(probabilities.shape[0], replace=False, p=probabilities)

        # Update used tokens list
        if token_id not in self._non_penalizable_tokens_ids:
            self._used_tokens_ids[sample_idx].append(token_id)

        return token_id

        Sample using individual priors for each sample_idx.
        Also updates individual priors for each sample in batch.
        We need individual priors to prevent the model from repeating the same tokens over and over again in one
        response.

        probabilities: Probabilities of each token. The distribution given by these probabilities
            is used by this function to sample the token.
        sample_idx : Integer between 0 and batch_size-1. We need it to figure out which token_log_prior line to use.
        temperature: Temperature for sampling. Temperature has to be a positive number.
        :return: Index of sampled token
        
cakechat/dialog_model/inference/candidates/sampling.py