def predict_log_prob_one_step(self,
                                  thought_vector,
                                  prev_hidden_states,
                                  prev_tokens_id,
                                  condition_id,
                                  temperature=1.0):
        """
        :param thought_vector:          shape == (batch_size, hidden_layer_dim), float32
        :param prev_hidden_states:      shape == (batch_size, 2 * hidden_layer_dim), float32
        :param prev_tokens_id:          shape == (batch_size, 1), int32
        :param condition_id:            shape == (batch_size, 1), int32
        :param temperature:             float32
        :return:
            updated_hidden_states:      shape == (batch_size, 2 * hidden_layer_dim), float32
            transformed_token_prob:     shape == (batch_size, vocab_size), float32
        """
        updated_hidden_states, token_prob = self.predict_prob_one_step(thought_vector, prev_hidden_states,
                                                                       prev_tokens_id, condition_id, temperature)

        token_log_prob = np.log(token_prob)
        return updated_hidden_states, token_log_prob

        :param thought_vector:          shape == (batch_size, hidden_layer_dim), float32
        :param prev_hidden_states:      shape == (batch_size, 2 * hidden_layer_dim), float32
        :param prev_tokens_id:          shape == (batch_size, 1), int32
        :param condition_id:            shape == (batch_size, 1), int32
        :param temperature:             float32
        :return:
            updated_hidden_states:      shape == (batch_size, 2 * hidden_layer_dim), float32
            transformed_token_prob:     shape == (batch_size, vocab_size), float32
        
cakechat/dialog_model/model.py