Function _body

examples/tensorflow/decoder/utils/decoding.py:247–285

(word_ids, cum_log_probs, finished, step, outputs, my_cache, extra_vars, op_self_cache, op_mem_cache)

            return tf.reduce_any(tf.logical_not(finished))

        def _body(word_ids, cum_log_probs, finished, step, outputs, my_cache, extra_vars, op_self_cache, op_mem_cache):
            logits, my_cache, op_self_cache, op_mem_cache = decoding_body(word_ids,
                                                                        step,
                                                                        extended_memory,
                                                                        extended_memory_sequence_length,
                                                                        my_cache,
                                                                        op_self_cache,
                                                                        op_mem_cache,
                                                                        embedding_table,
                                                                        decoding_args,
                                                                        decoder_type,
                                                                        extra_vars[1])

            end_ids = tf.fill([tf.shape(logits)[0]], decoding_args.end_id) # [batch_size * beam_width]
            eos_max_prob = tf.one_hot(end_ids, decoding_args.vocab_size,
                                      on_value=decoder_args.dtype.max,
                                      off_value=decoder_args.dtype.min)  # [batch_size * beam_width, vocab_size]
            
            # [batch_size * beam_width, vocab_size]
            logits = tf.where(finished, x=eos_max_prob, y=logits)
            logits = tf.cast(logits, tf.float32)
            
            output_id, next_cum_log_probs, finished, my_cache, \
                extra_vars, op_self_cache = search_word(beam_width,
                                                        decoding_args.vocab_size,
                                                        step,
                                                        logits,
                                                        cum_log_probs,
                                                        finished,
                                                        my_cache,
                                                        extra_vars,
                                                        op_self_cache,
                                                        search_method=search_method)
            cum_log_probs = tf.where(finished, x=cum_log_probs, y=next_cum_log_probs)
            
            outputs = outputs.write(step, output_id)
            finished = tf.logical_or(finished, tf.equal(output_id, decoding_args.end_id))

            return output_id, cum_log_probs, finished, step + 1, outputs, my_cache, extra_vars, op_self_cache, op_mem_cache

        # initialization
        batchxbeam = tf.shape(extended_memory)[0]