code: def beam_search_step(
      state, logits, eos_id, beam_width, is_first_step, length_penalty
  ):
      """logits.shape == [B*W, V]"""
      _, vocab_size = logits.shape
  
      bsz, beam_width = state.log_probs.shape
      onehot_eos = paddle.cast(
          nn.functional.one_hot(
              paddle.ones([1], "int64") * eos_id, vocab_size
          ),
          "int64",
      )  # [1, V]
  
      probs = paddle.log(nn.functional.softmax(logits))  # [B*W, V]
      probs = mask_prob(probs, onehot_eos, state.finished)  # [B*W, V]
      allprobs = (
          paddle.reshape(state.log_probs, [-1, 1]) + probs
      )  # [B*W, V]
  
      not_finished = 1 - paddle.reshape(
          state.finished, [-1, 1]
      )  # [B*W,1]
      not_eos = 1 - onehot_eos
      length_to_add = not_finished * not_eos  # [B*W,V]
      alllen = paddle.reshape(state.lengths, [-1, 1]) + length_to_add
  
      allprobs = paddle.reshape(allprobs, [-1, beam_width * vocab_size])
      alllen = paddle.reshape(alllen, [-1, beam_width * vocab_size])
      allscore = hyp_score(allprobs, alllen, length_penalty)
      if is_first_step:
          allscore = paddle.reshape(allscore, [bsz, beam_width, -1])[
              :, 0, :
          ]  # first step only consiter beam 0
      scores, idx = paddle.topk(allscore, k=beam_width)  # [B, W]
      next_beam_id = idx // vocab_size  # [B, W]
      next_word_id = idx % vocab_size
  
      gather_idx = paddle.concat(
          [
              paddle.nonzero(idx != -1)[:, :1],
              paddle.reshape(idx, [-1, 1]),
          ],
          1,
      )
      next_probs = paddle.reshape(
          paddle.gather_nd(allprobs, gather_idx), idx.shape
      )
      next_len = paddle.reshape(
          paddle.gather_nd(alllen, gather_idx), idx.shape
      )
  
      gather_idx = paddle.concat(
          [
              paddle.nonzero(next_beam_id != -1)[:, :1],
              paddle.reshape(next_beam_id, [-1, 1]),
          ],
          1,
      )
      next_finished = paddle.reshape(
          paddle.gather_nd(state.finished, gather_idx),
          state.finished.shape,
      )  # [gather new beam state according to new beam id]
  
      next_finished += paddle.cast(next_word_id == eos_id, "int64")
      next_finished = paddle.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
  
imports: from __future__ import division
  from __future__ import absolute_import
  from __future__ import print_function
  from __future__ import unicode_literals
  import sys
  import re
  import argparse
  import logging
  import json
  import numpy
  from collections import namedtuple
  import paddle
  import paddle.nn
  import numpy
  from paddlenlp.utils.log import logger
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