code: def _beam_search_generate(
      decoder: Seq2SeqDecoder,
      tokens=None,
      state=None,
      max_length=20,
      max_len_a=0.0,
      num_beams=4,
      temperature=1.0,
      top_k=50,
      top_p=1.0,
      bos_token_id=None,
      eos_token_id=None,
      do_sample=True,
      repetition_penalty=1.0,
      length_penalty=None,
      pad_token_id=0,
  ) -> torch.LongTensor:
      # 进行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
      )  # 这里要传入的是整个句子的长度
      if _eos_token_id != -1:  # 防止第一个位置为结束
          scores[:, _eos_token_id] = -1e12
      vocab_size = scores.size(1)
      assert (
          vocab_size >= num_beams
      ), "num_beams should be smaller than the number of vocabulary size."
  
      if do_sample:
          probs = F.softmax(scores, dim=-1) + 1e-12
          next_tokens = torch.multinomial(
              probs, num_samples=num_beams
          )  # (batch_size, num_beams)
          logits = probs.log()
          next_scores = logits.gather(
              dim=1, index=next_tokens
          )  # (batch_size, num_beams)
      else:
          scores = F.log_softmax(
              scores, dim=-1
          )  # (batch_size, vocab_size)
          # 得到(batch_size, num_beams), (batch_size, num_beams)
          next_scores, next_tokens = torch.topk(
              scores, num_beams, dim=1, largest=True, sorted=True
          )
  
      # 根据index来做顺序的调转
      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
      # 记录生成好的token (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
  
      #  用来记录已经生成好的token的长度
      cur_len = token_ids.size(1)
  
      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(
                  (tokens.size(0),),
                  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(
                  (tokens.size(0),),
                  fill_value=max_length,
                  dtype=torch.long,
              )
      hypos = [
          BeamHypotheses(
              num_beams,
              real_max_length,
              length_penalty,
              early_stopping=False,
          )
          for _ in range(batch_size)
      ]
      # 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]
              # 如果已经达到最大长度，就把eos的分数加大
              scores[:, _eos_token_id] = torch.where(
                  max_len_eos_mask, eos_scores + 1e32, eos_scores
              )
  
          if do_sample:
              if temperature > 0 and temperature != 1:
                  scores = scores / temperature
  
              # 多召回一个防止eos
              scores = top_k_top_p_filtering(
                  scores, top_k, top_p, min_tokens_to_keep=num_beams + 1
              )
              # 加上1e-12是为了避免https://github.com/pytorch/pytorch/pull/27523
              probs = F.softmax(scores, dim=-1) + 1e-12
  
              # 保证至少有一个不是eos的值
              _tokens = torch.multinomial(
                  probs, num_samples=num_beams + 1
              )  # batch_size' x (num_beams+1)
  
              logits = probs.log()
              # 防止全是这个beam的被选中了，且需要考虑eos被选择的情况
              _scores = logits.gather(
                  dim=1, index=_tokens
              )  # batch_size' x (num_beams+1)
              _scores = (
                  _scores + beam_scores[:, None]
              )  # batch_size' x (num_beams+1)
              # 从这里面再选择top的2*num_beam个
              _scores = _scores.view(
                  batch_size, num_beams * (num_beams + 1)
              )
              next_scores, ids = _scores.topk(
                  2 * num_beams, dim=1, largest=True, sorted=True
              )
              _tokens = _tokens.view(
                  batch_size, num_beams * (num_beams + 1)
              )
              next_tokens = _tokens.gather(
                  dim=1, index=ids
              )  # (batch_size, 2*num_beams)
              from_which_beam = torch.floor(
                  ids.float() / (num_beams + 1)
              ).long()  # (batch_size, 2*num_beams)
          else:
              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)
              next_scores, ids = torch.topk(
                  _scores,
                  2 * num_beams,
                  dim=1,
                  largest=True,
                  sorted=True,
              )  # (bsz, 2*num_beams)
              from_which_beam = torch.floor(
                  ids.float() / vocab_size
              ).long()  # (batch_size, 2*num_beams)
              next_tokens = (
                  ids % vocab_size
              )  # (batch_size, 2*num_beams)
  
          #  接下来需要组装下一个batch的结果。
          #  需要选定哪些留下来
          # next_scores, sorted_inds = next_scores.sort(dim=-1, descending=True)
          # next_tokens = next_tokens.gather(dim=1, index=sorted_inds)
          # from_which_beam = from_which_beam.gather(dim=1, index=sorted_inds)
  
          not_eos_mask = next_tokens.ne(_eos_token_id)  # 为1的地方不是eos
          keep_mask = not_eos_mask.cumsum(dim=1).le(
              num_beams
          )  # 为1的地方需要保留
          keep_mask = not_eos_mask.__and__(
              keep_mask
          )  # 为1的地方是需要进行下一步search的
  
          _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
          )  # 上面的token是来自哪个beam
          _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)
              )  # 表示的是indice
              eos_beam_idx = from_which_beam[:, :num_beams].reshape(
                  -1
              )  # 表示的是从哪个beam获取得到的
          else:
              # 将每个batch中在num_beam内的序列添加到结束中, 为1的地方需要结束了
              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)
                  # 是由于from_which_beam是 (batch_size, 2*num_beams)的，所以需要2*num_beams
                  eos_beam_idx = (
                      eos_batch_idx * num_beams * 2 + eos_beam_ind
                  )
                  eos_beam_idx = from_which_beam.view(-1)[
                      eos_beam_idx
                  ]  # 获取真实的从哪个beam获取的eos
              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()
                      # 之后需要在结尾新增一个eos
                      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,
                          )
  
          # 更改state状态, 重组token_ids
          reorder_inds = (
              batch_inds_with_numbeams_interval + _from_which_beam
          ).view(
              -1
          )  # flatten成一维
          state.reorder_state(reorder_inds)
          # 重新组织token_ids的状态
          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]
          # 把上面替换为非eos的词替换回eos
          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
  
file_path: fastNLP/fastNLP/modules/torch/generator/seq2seq_generator.py
imports: import torch
  from torch import nn
  import torch.nn.functional
  from ..decoder.seq2seq_decoder import Seq2SeqDecoder, State
  from functools import partial
label_desc: Use image node_id 5 for calling a global function with example usage: _beam_search_generate(decoder, tokens, state, max_length, max_len_a, num_beams, temperature, top_k, top_p, bos_token_id, eos_token_id, do_sample, repetition_penalty, length_penalty, pad_token_id) and returns: decoded