def predict_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
        """
        temperature = np.full_like(prev_tokens_id, temperature, dtype=np.float32)
        # shape == (batch_size, 1)

        token_prob, updated_hidden_states = self._models['decoder'].predict(
            [thought_vector, prev_tokens_id, condition_id, prev_hidden_states, temperature])
        return updated_hidden_states, token_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