code: def beam_search_step(
      state, logits, eos_id, beam_width, is_first_step, length_penalty
  ):
      """logits.shape == [B*W, V]"""
      (
          beam_size,
          vocab_size,
      ) = (
          logits.shape
      )  # as batch size=1 in this hub module. the first dim means bsz * beam_size equals beam_size
      logits_np = logits.numpy()
      for i in range(beam_size):
          logits_np[i][17963] = 0  # make [UNK] prob = 0
      logits = D.to_variable(logits_np)
  
      bsz, beam_width = state.log_probs.shape
      onehot_eos = L.cast(
          F.one_hot(L.ones([1], "int64") * eos_id, vocab_size), "int64"
      )  # [1, V]
  
      probs = L.log(L.softmax(logits))  # [B*W, V]
      probs = mask_prob(probs, onehot_eos, state.finished)  # [B*W, V]
      allprobs = L.reshape(state.log_probs, [-1, 1]) + probs  # [B*W, V]
  
      not_finished = 1 - L.reshape(state.finished, [-1, 1])  # [B*W,1]
      not_eos = 1 - onehot_eos
      length_to_add = not_finished * not_eos  # [B*W,V]
      alllen = L.reshape(state.lengths, [-1, 1]) + length_to_add
  
      allprobs = L.reshape(allprobs, [-1, beam_width * vocab_size])
      alllen = L.reshape(alllen, [-1, beam_width * vocab_size])
      allscore = hyp_score(allprobs, alllen, length_penalty)
      if is_first_step:
          allscore = L.reshape(allscore, [bsz, beam_width, -1])[
              :, 0, :
          ]  # first step only consiter beam 0
      scores, idx = L.topk(allscore, k=beam_width)  # [B, W]
      next_beam_id = idx // vocab_size  # [B, W]
      next_word_id = idx % vocab_size
  
      gather_idx = L.concat(
          [L.where(idx != -1)[:, :1], L.reshape(idx, [-1, 1])], 1
      )
      next_probs = L.reshape(
          L.gather_nd(allprobs, gather_idx), idx.shape
      )
      next_len = L.reshape(L.gather_nd(alllen, gather_idx), idx.shape)
  
      gather_idx = L.concat(
          [
              L.where(next_beam_id != -1)[:, :1],
              L.reshape(next_beam_id, [-1, 1]),
          ],
          1,
      )
      next_finished = L.reshape(
          L.gather_nd(state.finished, gather_idx), state.finished.shape
      )  # [gather new beam state according to new beam id]
  
      next_finished += L.cast(next_word_id == eos_id, "int64")
      next_finished = L.cast(next_finished > 0, "int64")
  
      next_state = BeamSearchState(
          log_probs=next_probs, lengths=next_len, finished=next_finished
      )
      output = BeamSearchOutput(
          scores=scores,
          predicted_ids=next_word_id,
          beam_parent_ids=next_beam_id,
      )
  
      return output, next_state
  
file_path: PaddleHub/modules/text/text_generation/reading_pictures_writing_poems_for_midautumn/MidAutumnPoetry/model/decode.py
imports: import re
  import numpy
  from collections import namedtuple
  import paddle.fluid
  import paddle.fluid.layers
  import paddle.fluid.dygraph
label_desc: Use image node_id 6 for calling a global function with example usage: beam_search_step(state, logits, eos_id, beam_width, is_first_step, length_penalty) and returns: output,  next_state