code: def _beam_search_generate(
      decoder: PromptBartDecoder,
      tokens=None,
      state=None,
      max_length=20,
      max_len_a=0.0,
      num_beams=4,
      bos_token_id=None,
      eos_token_id=None,
      do_sample=True,
      repetition_penalty=1.0,
      length_penalty=None,
      pad_token_id=0,
      restricter=None,
  ) -> torch.LongTensor:
      assert do_sample is False
      # beam search
      device = get_model_device(decoder)
      if tokens is None:
          if bos_token_id is None:
              raise RuntimeError(
                  "You have to specify either `tokens` or `bos_token_id`."
              )
          batch_size = state.num_samples()
          if batch_size is None:
              raise RuntimeError(
                  "Cannot infer the number of samples from `state`."
              )
          tokens = torch.full(
              [batch_size, 1], fill_value=bos_token_id, dtype=torch.long
          ).to(device)
      batch_size = tokens.size(0)
      if state.num_samples:
          assert (
              state.num_samples() == batch_size
          ), "The number of samples in `tokens` and `state` should match."
  
      if eos_token_id is None:
          _eos_token_id = -1
      else:
          _eos_token_id = eos_token_id
  
      scores = decoder.decode(tokens=tokens, state=state)
      vocab_size = scores.size(1)
      assert (
          vocab_size >= num_beams
      ), "num_beams should be smaller than the number of vocabulary size."
  
      scores = F.log_softmax(scores, dim=-1)  # (batch_size, vocab_size)
      if restricter is not None:
          _next_scores, _next_tokens = restricter(
              state, tokens, scores, num_beams + 1
          )
      else:
          # bsz x (num_beams+1)
          _next_scores, _next_tokens = torch.topk(
              scores, num_beams + 1, dim=1, largest=True, sorted=True
          )
  
      indices = torch.arange(batch_size, dtype=torch.long).to(device)
      indices = indices.repeat_interleave(num_beams)
      state.reorder_state(indices)
      tokens = tokens.index_select(
          dim=0, index=indices
      )  # batch_size * num_beams x length
  
      if max_len_a != 0:
          # (bsz x num_beams, )
          if state.encoder_mask is not None:
              max_lengths = (
                  state.encoder_mask.sum(dim=1).float() * max_len_a
              ).long() + max_length
          else:
              max_lengths = tokens.new_full(
                  (batch_size * num_beams,),
                  fill_value=max_length,
                  dtype=torch.long,
              )
          real_max_length = max_lengths.max().item()
      else:
          real_max_length = max_length
          if state.encoder_mask is not None:
              max_lengths = (
                  state.encoder_mask.new_ones(
                      state.encoder_mask.size(0)
                  ).long()
                  * max_length
              )
          else:
              max_lengths = tokens.new_full(
                  (batch_size * num_beams,),
                  fill_value=max_length,
                  dtype=torch.long,
              )
      hypos = [
          BeamHypotheses(
              num_beams,
              real_max_length,
              length_penalty,
              early_stopping=False,
          )
          for _ in range(batch_size)
      ]
  
      not_eos_mask = _next_tokens.ne(_eos_token_id)
      keep_mask = not_eos_mask.cumsum(dim=1).le(num_beams)
      keep_mask = not_eos_mask.__and__(keep_mask)
  
      next_tokens = _next_tokens.masked_select(keep_mask).view(
          batch_size, num_beams
      )
      next_scores = _next_scores.masked_select(keep_mask).view(
          batch_size, num_beams
      )
  
      rows, cols = not_eos_mask.eq(0)[:, :num_beams].nonzero(
          as_tuple=True
      )
  
      if len(rows) > 0:
          for row, col in zip(rows.tolist(), cols.tolist()):
              _token = torch.cat(
                  [
                      tokens[row * num_beams],
                      _next_tokens[row, col : col + 1],
                  ],
                  dim=0,
              )
              hypos[row].add(
                  _token.clone(), _next_scores[row, col].item()
              )
  
      # (batch_size, cur_len)
      token_ids = torch.cat([tokens, next_tokens.view(-1, 1)], dim=-1)
      dones = [False] * batch_size
  
      beam_scores = next_scores.view(-1)  # batch_size * num_beams
  
      cur_len = token_ids.size(1)
  
      # 0, num_beams, 2*num_beams, ...
      batch_inds_with_numbeams_interval = (
          (torch.arange(batch_size) * num_beams)
          .view(-1, 1)
          .to(token_ids)
      )
  
      while cur_len < real_max_length:
          scores = decoder.decode(
              token_ids, state
          )  # (bsz x num_beams, vocab_size)
          if repetition_penalty != 1.0:
              token_scores = scores.gather(dim=1, index=token_ids)
              lt_zero_mask = token_scores.lt(0).float()
              ge_zero_mask = lt_zero_mask.eq(0).float()
              token_scores = (
                  lt_zero_mask * repetition_penalty * token_scores
                  + ge_zero_mask / repetition_penalty * token_scores
              )
              scores.scatter_(dim=1, index=token_ids, src=token_scores)
  
          if _eos_token_id != -1:
              max_len_eos_mask = max_lengths.eq(cur_len + 1)
              eos_scores = scores[:, _eos_token_id]
              scores[:, _eos_token_id] = torch.where(
                  max_len_eos_mask, eos_scores + 1e32, eos_scores
              )
  
          scores = F.log_softmax(
              scores, dim=-1
          )  # (batch_size * num_beams, vocab_size)
          _scores = (
              scores + beam_scores[:, None]
          )  # (batch_size * num_beams, vocab_size)
          _scores = _scores.view(
              batch_size, -1
          )  # (batch_size, num_beams*vocab_size)
  
          if restricter is not None:
              next_scores, ids = restricter(
                  state, token_ids, _scores, 2 * num_beams
              )
          else:
              next_scores, ids = torch.topk(
                  _scores,
                  2 * num_beams,
                  dim=1,
                  largest=True,
                  sorted=True,
              )  # (bsz, 2*num_beams)
          from_which_beam = (
              ids // vocab_size
          )  # (batch_size, 2*num_beams)
          next_tokens = ids % vocab_size  # (batch_size, 2*num_beams)
  
          not_eos_mask = next_tokens.ne(_eos_token_id)
          keep_mask = not_eos_mask.cumsum(dim=1).le(num_beams)
          keep_mask = not_eos_mask.__and__(keep_mask)
  
          _next_tokens = next_tokens.masked_select(keep_mask).view(
              -1, 1
          )
          _from_which_beam = from_which_beam.masked_select(
              keep_mask
          ).view(batch_size, num_beams)
          _next_scores = next_scores.masked_select(keep_mask).view(
              batch_size, num_beams
          )
          beam_scores = _next_scores.view(-1)
  
          flag = True
          if cur_len + 1 == real_max_length:
              eos_batch_idx = (
                  torch.arange(batch_size)
                  .to(next_tokens)
                  .repeat_interleave(repeats=num_beams, dim=0)
              )
              eos_beam_ind = (
                  torch.arange(num_beams)
                  .to(token_ids)
                  .repeat(batch_size)
              )
              eos_beam_idx = from_which_beam[:, :num_beams].reshape(-1)
          else:
              effective_eos_mask = next_tokens[:, :num_beams].eq(
                  _eos_token_id
              )  # batch_size x num_beams
              if effective_eos_mask.sum().gt(0):
                  (
                      eos_batch_idx,
                      eos_beam_ind,
                  ) = effective_eos_mask.nonzero(as_tuple=True)
                  eos_beam_idx = (
                      eos_batch_idx * num_beams * 2 + eos_beam_ind
                  )
                  eos_beam_idx = from_which_beam.view(-1)[eos_beam_idx]
              else:
                  flag = False
  
          if flag:
              _token_ids = torch.cat([token_ids, _next_tokens], dim=-1)
              for batch_idx, beam_ind, beam_idx in zip(
                  eos_batch_idx.tolist(),
                  eos_beam_ind.tolist(),
                  eos_beam_idx.tolist(),
              ):
                  if not dones[batch_idx]:
                      score = next_scores[batch_idx, beam_ind].item()
                      if _eos_token_id != -1:
                          hypos[batch_idx].add(
                              _token_ids[
                                  batch_idx * num_beams + beam_idx,
                                  :cur_len,
                              ].clone(),
                              score,
                          )
                      else:
                          hypos[batch_idx].add(
                              _token_ids[
                                  batch_idx * num_beams + beam_idx
                              ].clone(),
                              score,
                          )
  
          reorder_inds = (
              batch_inds_with_numbeams_interval + _from_which_beam
          ).view(
              -1
          )  # flatten
          state.reorder_state(reorder_inds)
          token_ids = torch.cat(
              [
                  token_ids.index_select(index=reorder_inds, dim=0),
                  _next_tokens,
              ],
              dim=-1,
          )
  
          for batch_idx in range(batch_size):
              dones[batch_idx] = (
                  dones[batch_idx]
                  or hypos[batch_idx].is_done(
                      next_scores[batch_idx, 0].item()
                  )
                  or max_lengths[batch_idx * num_beams] == cur_len + 1
              )
  
          cur_len += 1
  
          if all(dones):
              break
  
      # select the best hypotheses
      tgt_len = token_ids.new_zeros(batch_size)
      best = []
  
      for i, hypotheses in enumerate(hypos):
          best_hyp = max(hypotheses.hyp, key=lambda x: x[0])[1]
          if _eos_token_id != -1:
              best_hyp = torch.cat(
                  [best_hyp, best_hyp.new_ones(1) * _eos_token_id]
              )
          tgt_len[i] = len(best_hyp)
          best.append(best_hyp)
  
      # generate target batch
      decoded = token_ids.new_zeros(
          batch_size, tgt_len.max().item()
      ).fill_(pad_token_id)
      for i, hypo in enumerate(best):
          decoded[i, : tgt_len[i]] = hypo
  
      return decoded
  
imports: import torch
  from torch import nn
  from torch.nn import functional
  from transformers.models.bart.configuration_bart import BartConfig
  from .modeling_bart import BartModel
  from ..utils.util import avg_token_embeddings, seq_to_mask, get_model_device
  from functools import partial
  from typing import Union
label_desc: Use image node_id 3 for calling a global function with example usage: _beam_search_generate(decoder, tokens, state, max_length, max_len_a, num_beams, bos_token_id, eos_token_id, do_sample, repetition_penalty, length_penalty, pad_token_id, restricter) and returns: decoded