docstr: Beam search
  Args:
      model_inputs (dict): input ids.
      beam_size (int, optional, defaults to 3): beam size of beam search.
      generate_length (int, optional, defaults to 100): maximum generation length.
      repetition_penalty (float, optional, defaults to 1.0): repetition penalty coefficient, 1.0 means no penalty.
      repetition_window (int, optional, defaults to None): window size of repetition penalty, None means that all output tokens are penalized.
docstr_tok: ["Beam","search","Args",":","model_inputs","(","dict",")",":","input","ids",".","beam_size","(","int",",","optional",",","defaults","to","3",")",":","beam","size","of","beam","search",".","generate_length","(","int",",","optional",",","defaults","to","100",")",":","maximum","generation","length",".","repetition_penalty","(","float",",","optional",",","defaults","to","1.0",")",":","repetition","penalty","coefficient",",","1.0","means","no","penalty",".","repetition_window","(","int",",","optional",",","defaults","to","None",")",":","window","size","of","repetition","penalty",",","None","means","that","all","output","tokens","are","penalized","."]
code: def _decode(
      self,
      model_inputs,
      other_info,
      beam_size=3,
      max_length=100,
      repetition_penalty=1.0,
      repetition_window=None,
  ):
      """
      Beam search
      Args:
          model_inputs (dict): input ids.
          beam_size (int, optional, defaults to 3): beam size of beam search.
          generate_length (int, optional, defaults to 100): maximum generation length.
          repetition_penalty (float, optional, defaults to 1.0): repetition penalty coefficient, 1.0 means no penalty.
          repetition_window (int, optional, defaults to None): window size of repetition penalty, None means that all output tokens are penalized.
      """  # noqa: E501
      # generate_length + 1 for EOS token
      max_length += 1
  
      # expand dimmension
      batch_size = model_inputs["input"].size(0)
      input: torch.Tensor = (
          model_inputs["input"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      input_sub: torch.Tensor = (
          model_inputs["input_sub"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      input_pos: torch.Tensor = (
          model_inputs["input_pos"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      context: torch.Tensor = (
          model_inputs["context"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      sample_ids: torch.Tensor = (
          model_inputs["sample_idx"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      num_segments: torch.Tensor = (
          model_inputs["num_segments"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      segment: torch.Tensor = (
          model_inputs["segment"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      segment_rel_offset: torch.Tensor = (
          model_inputs["segment_rel_offset"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      segment_rel: torch.Tensor = (
          model_inputs["segment_rel"]
          .unsqueeze(1)
          .expand(batch_size, beam_size, -1)
          .contiguous()
          .view(batch_size * beam_size, -1)
      )
      ext_table_ids: torch.Tensor = model_inputs["batch_ext_table_ids"]
      ext_table_sub: torch.Tensor = model_inputs["batch_ext_table_sub"]
      ext_table_ids_cpu = ext_table_ids.cpu()
      ext_table_sub_cpu = ext_table_sub.cpu()
  
      done = [False for _ in range(batch_size)]
  
      beam_scores = torch.zeros(
          (batch_size, beam_size),
          dtype=torch.float,
          device=input.device,
      )
      beam_scores[:, 1:] = -1e9
      beam_scores = beam_scores.view(-1)
  
      # generated hypotheses
      generated_hyps = [
          BeamHypotheses(
              beam_size,
              max_length,
              length_penalty=1,
              early_stopping=False,
          )
          for _ in range(batch_size)
      ]
  
      pred_start_index = input.size(-1)
      _, _, past_key_values = self.model.inference(
          input=input,
          input_sub=input_sub,
          position=input_pos,
          context=context,
          sample_ids=sample_ids,
          num_segments=num_segments,
          segment=segment,
          segment_rel_offset=segment_rel_offset,
          segment_rel=segment_rel,
          ext_table_ids=ext_table_ids,
          ext_table_sub=ext_table_sub,
          past_key_values=None,
      )
  
      beam_states = []
      for sent_id in range(batch_size):
          instance_beam_states = []
  
          for beam_id in range(beam_size):
              instance_beam_states.append(
                  {
                      "idx": 0,
                      "ans": [],
                      "nx_token_id": self.tokenizer.bos_id,
                      "nx_token_sub": 0,
                      "nx_segment_id": other_info[sent_id][
                          "predict_segments"
                      ][0][0],
                      "nx_position": 0,
                  }
              )
          beam_states.append(instance_beam_states)
      for i in range(max_length + 1):
          tmp_input = []
          tmp_input_sub = []
          tmp_position = []
          tmp_segment = []
          for sent_id in range(batch_size):
              for beam_id in range(beam_size):
                  tmp_input.append(
                      beam_states[sent_id][beam_id]["nx_token_id"]
                  )
                  tmp_input_sub.append(
                      beam_states[sent_id][beam_id]["nx_token_sub"]
                  )
                  tmp_position.append(
                      beam_states[sent_id][beam_id]["nx_position"]
                  )
                  tmp_segment.append(
                      beam_states[sent_id][beam_id]["nx_segment_id"]
                  )
          with torch.no_grad():
              input = torch.cat(
                  [
                      input,
                      torch.tensor(
                          tmp_input,
                          dtype=torch.int32,
                          device=input.device,
                      ).view(batch_size * beam_size, 1),
                  ],
                  dim=-1,
              )
              logits, _, past_key_values = self.model.inference(
                  input=input[:, -1:],
                  input_sub=torch.tensor(
                      tmp_input_sub,
                      dtype=torch.int32,
                      device=input.device,
                  ).view(batch_size * beam_size, 1),
                  position=torch.tensor(
                      tmp_position,
                      dtype=torch.int32,
                      device=input.device,
                  ).view(batch_size * beam_size, 1),
                  context=torch.ones(
                      batch_size * beam_size,
                      dtype=torch.bool,
                      device=input.device,
                  ).view(batch_size * beam_size, 1),
                  sample_ids=torch.zeros(
                      batch_size * beam_size,
                      dtype=torch.int32,
                      device=input.device,
                  ).view(batch_size * beam_size, 1),
                  num_segments=num_segments[:, -1:],
                  segment=torch.tensor(
                      tmp_segment,
                      dtype=torch.int32,
                      device=input.device,
                  ).view(batch_size * beam_size, 1),
                  segment_rel_offset=segment_rel_offset[:, -1:],
                  segment_rel=segment_rel,
                  ext_table_ids=ext_table_ids,
                  ext_table_sub=ext_table_sub,
                  past_key_values=past_key_values,
              )
              logits = logits[:, -1, :]
  
          # skip all steps when we are done with each sentence
          if all(done):
              break
  
          for sent_id in range(batch_size):
              if (
                  self.tokenizer.unk_id
                  not in other_info[sent_id]["ext_table"]
              ):
                  # unk is not allowed, mask unk
                  logits[
                      sent_id * beam_size : (sent_id + 1) * beam_size,
                      self.tokenizer.unk_id,
                  ] = -10000
              ext_ids = set()
              for v in other_info[sent_id]["ext_table"].keys():
                  ext_ids.add(v)
              for ext_id in range(
                  self.tokenizer.vocab_size,
                  self.tokenizer.vocab_size + ext_table_ids.size(0),
              ):
                  if ext_id not in ext_ids:
                      logits[
                          sent_id
                          * beam_size : (sent_id + 1)
                          * beam_size,
                          ext_id,
                      ] = -10000
  
          apply_repetition_penalty(
              logits,
              batch_size,
              beam_size,
              input,
              repetition_penalty,
              pred_start_index,
              input.size(-1) - 1,
              repetition_window,
          )
          scores = F.log_softmax(logits, dim=-1)
          next_scores = scores + beam_scores[:, None].expand_as(
              scores
          )  # (batch_size * beam_size, vocab_size)
  
          # re-organize to group the beam together (we are keeping top hypothesis accross beams)
          next_scores = next_scores.view(
              batch_size, -1
          )  # (batch_size, beam_size * vocab_size)
          next_scores, next_words = torch.topk(
              next_scores,
              2 * beam_size,
              dim=1,
              largest=True,
              sorted=True,
          )
          assert (
              next_scores.size()
              == next_words.size()
              == (batch_size, 2 * beam_size)
          )
          next_beam_states = []
  
          for sent_id in range(batch_size):
              # if we are done with this sentence
              done[sent_id] = done[sent_id] or generated_hyps[
                  sent_id
              ].is_done(next_scores[sent_id].max().item(), i)
              if done[sent_id]:
                  next_beam_states.append(
                      [
                          (
                              {
                                  "idx": 0,
                                  "ans": [],
                                  "nx_token_id": 0,
                                  "nx_token_sub": 0,
                                  "nx_segment_id": 0,
                                  "nx_position": 0,
                              },
                              0,
                              0,
                          )
                      ]
                      * beam_size
                  )  # pad the batch
                  continue
  
              # next sentence beam content
              next_instance_beam_states = []
  
              # next words for this sentence
              for idx, value in zip(
                  next_words[sent_id], next_scores[sent_id]
              ):
                  # get beam and word IDs
                  beam_id = torch.div(
                      idx, scores.size(-1), rounding_mode="floor"
                  ).item()
                  word_id = (idx % scores.size(-1)).item()
  
                  curr_info = beam_states[sent_id][beam_id]
                  # end of sentence, or next word
                  if (
                      word_id == self.tokenizer.eos_id
                      and (
                          curr_info["idx"] + 1
                          == len(
                              other_info[sent_id]["predict_segments"]
                          )
                      )
                  ) or i == max_length:
                      generated_hyps[sent_id].add(
                          beam_states[sent_id][beam_id]["ans"]
                          + [
                              (
                                  word_id,
                                  other_info[sent_id][
                                      "predict_segments"
                                  ][curr_info["idx"]][1],
                              )
                          ],
                          value.item(),
                      )
                  elif word_id == self.tokenizer.eos_id:
                      next_instance_beam_states.append(
                          (
                              {
                                  "idx": curr_info["idx"] + 1,
                                  "ans": curr_info["ans"]
                                  + [
                                      (
                                          word_id,
                                          other_info[sent_id][
                                              "predict_segments"
                                          ][curr_info["idx"]][1],
                                      )
                                  ],
                                  "nx_token_id": self.tokenizer.bos_id,
                                  "nx_token_sub": 0,
                                  "nx_segment_id": other_info[sent_id][
                                      "predict_segments"
                                  ][curr_info["idx"] + 1][0],
                                  "nx_position": 0,
                              },
                              value.item(),
                              sent_id * beam_size + beam_id,
                          )
                      )
  
                  else:
                      raw_word_id = word_id
                      word_id_sub = 0
                      if word_id >= self.tokenizer.vocab_size:
                          word_id -= self.tokenizer.vocab_size
                          word_id_sub = int(
                              ext_table_sub_cpu[word_id].item()
                          )
                          word_id = int(
                              ext_table_ids_cpu[word_id].item()
                          )
  
                      next_instance_beam_states.append(
                          (
                              {
                                  "idx": curr_info["idx"],
                                  "ans": curr_info["ans"]
                                  + [
                                      (
                                          raw_word_id,
                                          other_info[sent_id][
                                              "predict_segments"
                                          ][curr_info["idx"]][1],
                                      )
                                  ],
                                  "nx_token_id": word_id,
                                  "nx_token_sub": word_id_sub,
                                  "nx_segment_id": curr_info[
                                      "nx_segment_id"
                                  ],
                                  "nx_position": curr_info[
                                      "nx_position"
                                  ]
                                  + 1,
                              },
                              value.item(),
                              sent_id * beam_size + beam_id,
                          )
                      )
  
                  # the beam for next step is full
                  if len(next_instance_beam_states) == beam_size:
                      break
  
              # update next beam content
              assert (
                  len(next_instance_beam_states) == 0
                  if i == max_length
                  else beam_size
              )
              next_beam_states.append(next_instance_beam_states)
  
          # we have reached the last step
          if i == max_length:
              break
  
          # sanity check / prepare next batch
          beam_reorder_idx = []
          beam_new_scores = []
          beam_states = []
          for sent_id in range(batch_size):
              instance_beam_states = []
              for beam_id in range(beam_size):
                  state, value, beam_idx = next_beam_states[sent_id][
                      beam_id
                  ]
                  beam_reorder_idx.append(beam_idx)
                  beam_new_scores.append(value)
                  instance_beam_states.append(state)
              beam_states.append(instance_beam_states)
  
          input = input[beam_reorder_idx, :]
          beam_scores = torch.tensor(
              beam_new_scores, dtype=torch.float, device=input.device
          )
          for kw in past_key_values.keys():
              if kw == "buffer":
                  buf_list = past_key_values[kw]
                  nw_buf_list = []
                  for buf in buf_list:
                      if buf is None:
                          nw_buf_list.append((None, None))
                      else:
                          k_buf, v_buf = buf
                          nw_buf_list.append(
                              (
                                  k_buf[beam_reorder_idx, :],
                                  v_buf[beam_reorder_idx, :],
                              )
                          )
                  past_key_values[kw] = nw_buf_list
              else:
                  past_key_values[kw] = past_key_values[kw][
                      beam_reorder_idx, :
                  ]
  
      # select the best hypotheses
      results = []
      for sent_id, hypotheses in enumerate(generated_hyps):
          best_hyp = max(hypotheses.hyp, key=lambda x: x[0])[1]
          results.append(best_hyp)
      return results
  
file_path: DeepKE/example/llm/CPM-Bee/src/cpm_live/generation/bee.py
imports: from typing import Any, Dict, List, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from .generation_utils import BeamHypotheses, apply_repetition_penalty
  from ..tokenizers.bee import CPMBeeTokenizer
  from ..models.bee import CPMBee
  from ..training_tasks.bee.pretrain import convert_data_to_id
  from ..utils import pad
label_desc: Use image node_id 1 for calling the CPMBeeBeamSearch obj's underlying member method code with example usage: obj._decode(model_inputs, other_info, beam_size, max_length, repetition_penalty, repetition_window) and returns: results