docstr: # Args.
      input_token_ids: IntTensor, (batch_size, max_input_length),
          input hidden state to decoder.
      input_lengths: IntTensor, (batch_size),
          the lengths of input context sequences.
      gen_length: int, the number of tokens to generate.
      local_batch_size: int, optional, a batch size of
          local iteration. (disabled)
      eos_token_id: int, eos token id.
      beam_width: int, number of beams for beam search.
          If 1, sampling decode will be used.
      top_k: IntTensor, (batch_size,) top-k sampling.
          The number of most probable tokens to keep
          for sampling per sentence in a batcch.
      top_p: FloatTensor, (batch_size,), top-p sampling.
          The cumulative probability
          of to filter the set of most probable tokens.
      top_p_decay: FloatTensor, (batch_size,)
          The decay of top-p value for top_p sampling.
      top_p_min: FloatTensor, (batch_size,)
          The minimum top p values in top-p decaying.
      top_p_reset_ids: IntTensor, (batch_size,)
          reset ids for resetting top_p values for top p sampling
      temperature: FloatTensor, (batch_size,),
          The temperature value for smoothing the logit distribution.
      repetition_penalty: FloatTensor, (batch_size,),
          The repetition penalty.
      presence_penalty: FloatTensor, (batch_size,),
          The presence penalty, which is exclusive with
          repetition_penalty.
          Only one of repetition and presence penalties is allowed.
      min_length: IntTensor, (batch_size,),
          Minimum length for each sentences. EOS is masked if length is
          below min.
      len_penalty: FloatTensor, (batch_size,)
          The exponent of the length penalty of beam scores.
      beam_search_diversity_rate: FloatTensor, (batch_size,),
          The diversity rate of beam search.
      stop_words_list: IntTensor, (batch_size, 2, stop_words_length)
          When FT generates words in this list, it will stop the
          generation. An extension of stop id.
      bad_words_list IntTensor, (batch_size, 2, bad_words_length)
          The words in the list will never be sampled.
      sequence_limit_lengths: IntTensor, (batch_size,), The maximum
          length of a generated sequence.
      memory_length: int, the length of cache memory. If None, it will
          be max_input_length + gen_length.
  # Returns
      IntTensor, (batch_size, beam_width, max_seq_length) output
      token ids.
docstr_tok: ["#","Args",".","input_token_ids",":","IntTensor",",","(","batch_size",",","max_input_length",")",",","input","hidden","state","to","decoder",".","input_lengths",":","IntTensor",",","(","batch_size",")",",","the","lengths","of","input","context","sequences",".","gen_length",":","int",",","the","number","of","tokens","to","generate",".","local_batch_size",":","int",",","optional",",","a","batch","size","of","local","iteration",".","(","disabled",")","eos_token_id",":","int",",","eos","token","id",".","beam_width",":","int",",","number","of","beams","for","beam","search",".","If","1",",","sampling","decode","will","be","used",".","top_k",":","IntTensor",",","(","batch_size",",",")","top-k","sampling",".","The","number","of","most","probable","tokens","to","keep","for","sampling","per","sentence","in","a","batcch",".","top_p",":","FloatTensor",",","(","batch_size",",",")",",","top-p","sampling",".","The","cumulative","probability","of","to","filter","the","set","of","most","probable","tokens",".","top_p_decay",":","FloatTensor",",","(","batch_size",",",")","The","decay","of","top-p","value","for","top_p","sampling",".","top_p_min",":","FloatTensor",",","(","batch_size",",",")","The","minimum","top","p","values","in","top-p","decaying",".","top_p_reset_ids",":","IntTensor",",","(","batch_size",",",")","reset","ids","for","resetting","top_p","values","for","top","p","sampling","temperature",":","FloatTensor",",","(","batch_size",",",")",",","The","temperature","value","for","smoothing","the","logit","distribution",".","repetition_penalty",":","FloatTensor",",","(","batch_size",",",")",",","The","repetition","penalty",".","presence_penalty",":","FloatTensor",",","(","batch_size",",",")",",","The","presence","penalty",",","which","is","exclusive","with","repetition_penalty",".","Only","one","of","repetition","and","presence","penalties","is","allowed",".","min_length",":","IntTensor",",","(","batch_size",",",")",",","Minimum","length","for","each","sentences",".","EOS","is","masked","if","length","is","below","min",".","len_penalty",":","FloatTensor",",","(","batch_size",",",")","The","exponent","of","the","length","penalty","of","beam","scores",".","beam_search_diversity_rate",":","FloatTensor",",","(","batch_size",",",")",",","The","diversity","rate","of","beam","search",".","stop_words_list",":","IntTensor",",","(","batch_size",",","2",",","stop_words_length",")","When","FT","generates","words","in","this","list",",","it","will","stop","the","generation",".","An","extension","of","stop","id",".","bad_words_list","IntTensor",",","(","batch_size",",","2",",","bad_words_length",")","The","words","in","the","list","will","never","be","sampled",".","sequence_limit_lengths",":","IntTensor",",","(","batch_size",",",")",",","The","maximum","length","of","a","generated","sequence",".","memory_length",":","int",",","the","length","of","cache","memory",".","If","None",",","it","will","be","max_input_length","+","gen_length",".","#","Returns","IntTensor",",","(","batch_size",",","beam_width",",","max_seq_length",")","output","token","ids","."]
code: def generate(
      self,
      input_token_ids: torch.IntTensor,
      input_lengths: torch.IntTensor,
      gen_length: int,
      eos_token_id: Optional[int] = None,
      local_batch_size: Optional[int] = None,
      beam_width: int = 1,
      top_k: Optional[torch.IntTensor] = None,
      top_p: Optional[torch.FloatTensor] = None,
      top_p_decay: Optional[torch.FloatTensor] = None,
      top_p_min: Optional[torch.FloatTensor] = None,
      top_p_reset_ids: Optional[torch.IntTensor] = None,
      temperature: Optional[torch.FloatTensor] = None,
      repetition_penalty: Optional[torch.FloatTensor] = None,
      presence_penalty: Optional[torch.FloatTensor] = None,
      min_length: Optional[torch.IntTensor] = None,
      len_penalty: Optional[torch.FloatTensor] = None,
      beam_search_diversity_rate: Optional[torch.FloatTensor] = None,
      stop_words_list: Optional[torch.IntTensor] = None,
      bad_words_list: Optional[torch.IntTensor] = None,
      sequence_limit_lengths: Optional[torch.IntTensor] = None,
      random_seed: Optional[torch.LongTensor] = None,
      memory_length: Optional[int] = None,
      return_output_length: bool = False,
      return_log_probs: bool = False,
  ):
      """
  
      # Args.
          input_token_ids: IntTensor, (batch_size, max_input_length),
              input hidden state to decoder.
          input_lengths: IntTensor, (batch_size),
              the lengths of input context sequences.
          gen_length: int, the number of tokens to generate.
          local_batch_size: int, optional, a batch size of
              local iteration. (disabled)
          eos_token_id: int, eos token id.
          beam_width: int, number of beams for beam search.
              If 1, sampling decode will be used.
          top_k: IntTensor, (batch_size,) top-k sampling.
              The number of most probable tokens to keep
              for sampling per sentence in a batcch.
          top_p: FloatTensor, (batch_size,), top-p sampling.
              The cumulative probability
              of to filter the set of most probable tokens.
          top_p_decay: FloatTensor, (batch_size,)
              The decay of top-p value for top_p sampling.
          top_p_min: FloatTensor, (batch_size,)
              The minimum top p values in top-p decaying.
          top_p_reset_ids: IntTensor, (batch_size,)
              reset ids for resetting top_p values for top p sampling
          temperature: FloatTensor, (batch_size,),
              The temperature value for smoothing the logit distribution.
          repetition_penalty: FloatTensor, (batch_size,),
              The repetition penalty.
          presence_penalty: FloatTensor, (batch_size,),
              The presence penalty, which is exclusive with
              repetition_penalty.
              Only one of repetition and presence penalties is allowed.
          min_length: IntTensor, (batch_size,),
              Minimum length for each sentences. EOS is masked if length is
              below min.
          len_penalty: FloatTensor, (batch_size,)
              The exponent of the length penalty of beam scores.
          beam_search_diversity_rate: FloatTensor, (batch_size,),
              The diversity rate of beam search.
          stop_words_list: IntTensor, (batch_size, 2, stop_words_length)
              When FT generates words in this list, it will stop the
              generation. An extension of stop id.
          bad_words_list IntTensor, (batch_size, 2, bad_words_length)
              The words in the list will never be sampled.
          sequence_limit_lengths: IntTensor, (batch_size,), The maximum
              length of a generated sequence.
          memory_length: int, the length of cache memory. If None, it will
              be max_input_length + gen_length.
      # Returns
          IntTensor, (batch_size, beam_width, max_seq_length) output
          token ids.
      """
      assert (
          self.weight is not None
      ), "Please call load() first to initialize weights."
  
      input_token_ids = input_token_ids.type(torch.int32).to(
          self.device
      )
      input_lengths = input_lengths.type(torch.int32).to(self.device)
  
      batch_size = len(input_token_ids)
      max_input_length = input_token_ids.shape[-1]
      max_seq_length = max_input_length + gen_length
      memory_length = memory_length or max_seq_length
  
      # TODO: Enable local batch later. We currently disable local batching due to # noqa: E501
      #   an input mismatch issue of FT's decode_op: FT's decode_op requires logits # noqa: E501
      #   of shape (batch_size, ...) but we have logits of shape (local_batch_size, ...) # noqa: E501
      #   After fixing FT's side, we will enable local batch.
      # local_batch_size = local_batch_size or self.get_local_batch_size(batch_size) # noqa: E501
      # num_local_batches, last_chunk = divmod(batch_size, local_batch_size)
      # if last_chunk > 0:
      #     num_local_batches += 1
      assert local_batch_size is None or local_batch_size == batch_size
      local_batch_size = batch_size
      num_local_batches = 1
  
      device = self.device
  
      eos_token_id = (
          eos_token_id
          if eos_token_id is not None
          else self.config.end_id
      )
      assert (
          eos_token_id is not None
      ), "eos_token-id must be specified in generation."
      eos_token_ids = eos_token_id * torch.ones(
          batch_size, dtype=torch.int32, device=device
      )
      assert repetition_penalty is None or presence_penalty is None, (
          "Found ambiguous parameters repetition_penalty and "
          "presence_penalty which are mutually exclusive. "
          "Please provide one of repetition_penalty and presence_penalty."
      )
      # Setup decoder_op prior to calling the forward function.
      self.decode_op.setup(
          batch_size,
          beam_width,
          top_k,
          top_p,
          temperature,
          repetition_penalty,
          presence_penalty,
          min_length,
          len_penalty,
          beam_search_diversity_rate,
          random_seed,
          top_p_decay,
          top_p_min,
          top_p_reset_ids,
      )
  
      # Prepare input and output arguments.
      if beam_width > 1:
          # Tiling for beam search.
          input_token_ids = input_token_ids.repeat(1, beam_width).view(
              batch_size * beam_width, -1
          )
          input_lengths = (
              input_lengths.view(-1, 1).repeat(1, beam_width).view(-1)
          )
          if sequence_limit_lengths is not None:
              sequence_limit_lengths = (
                  sequence_limit_lengths.view(-1, 1)
                  .repeat(1, beam_width)
                  .view(-1)
              )
          # src/tgt cache indirections.
          cache_indirection = torch.zeros(
              (2, batch_size, beam_width, memory_length),
              dtype=torch.int32,
              device=device,
          )
          parent_ids = torch.zeros(
              max_seq_length,
              batch_size * beam_width,
              dtype=torch.int32,
              device=device,
          )
      else:
          cache_indirection = None
          src_cache_indirection = None
          tgt_cache_indirection = None
          parent_ids = None
  
      pad_lengths = max_input_length - input_lengths
      # Since tril() doesn't support bf16 dtype,
      # we create of bool type and then cast it to dtype.
      attention_mask = (
          torch.ones(
              (max_input_length, max_input_length),
              dtype=torch.bool,
              device=device,
          )
          .tril()
          .unsqueeze(0)
          .tile(input_token_ids.shape[0], 1, 1)
          .to(self.dtype)
      )
      for b, input_length in enumerate(input_lengths):
          attention_mask[b, input_length:, ...] = 0
      masked_tokens = self.generate_pad_mask(
          input_lengths, memory_length
      )
      finished = torch.zeros_like(input_lengths).bool()
      sequence_lengths = (max_input_length - 1) * torch.ones_like(
          input_lengths
      )
  
      if return_log_probs or beam_width > 1:
          cum_log_probs = torch.zeros(
              batch_size * beam_width, device=device
          )
          output_log_probs = torch.zeros(
              (gen_length, batch_size * beam_width), device=device
          )
      else:
          cum_log_probs = None
          output_log_probs = None
  
      # Contiguous buffer for each decode_op step,
      # it will be transposed tensor for the final output.
      output_token_ids = torch.zeros(
          (max_seq_length, batch_size * beam_width),
          dtype=torch.int32,
          device=device,
      )
      output_token_ids[:max_input_length, ...] = input_token_ids.T
  
      if comm.is_pipeline_group_first():
          # Prepare input tensors of decoder.
          input_embeds = self.word_embedding(input_token_ids)
          if self.position_encoding is not None:
              position_ids = torch.arange(
                  0, max_input_length, dtype=torch.int, device=device
              )
              position_ids = position_ids.unsqueeze(0).view(
                  -1, max_input_length
              )
              input_embeds += self.position_encoding(position_ids)
          if self.pre_decoder_layernorm is not None:
              input_embeds = self.pre_decoder_layernorm(input_embeds)
      else:
          # Dummy input_embeds
          input_embeds = torch.empty(
              size=(
                  batch_size * beam_width,
                  max_input_length,
                  self.context_decoder.hidden_size,
              ),
              dtype=self.context_decoder.dtype,
              device=device,
          )
  
      use_shared_contexts = (
          (self.shared_contexts_ratio > 0.0)
          and (max_input_length >= 1)
          and (batch_size > 1)
      )
      batch_to_compact, compact_to_batch = None, None
      if use_shared_contexts:
          find_context_duplications = (
              torch.ops.fastertransformer.find_context_duplications
          )
          (
              batch_to_compact,
              compact_to_batch,
          ) = find_context_duplications(input_token_ids)
          use_shared_contexts = (
              compact_to_batch.shape[0]
              <= self.shared_contexts_ratio * batch_size
          )
  
          if not use_shared_contexts:
              batch_to_compact, compact_to_batch = None, None
  
      profiler.start("ft-context-decoder")
      (
          _,
          k_cache,
          v_cache,
          last_token_hidden_states,
      ) = self.context_decoder.forward(
          input_embeds=input_embeds,
          attention_mask=attention_mask,
          input_lengths=input_lengths,
          memory_length=memory_length,
          batch_to_compact_index=batch_to_compact,
          compact_index=compact_to_batch,
      )
      profiler.stop("ft-context-decoder")
  
      for step in range(max_input_length, max_seq_length):
          src_indir_idx = (step - max_input_length) % 2
          tgt_indir_idx = 1 - src_indir_idx
  
          is_generation_done = torch.tensor(
              [True], dtype=torch.bool, device=device
          )
          for ite in range(num_local_batches):
              # The indices of the current local batch-beam.
              bbidx = range(
                  ite * local_batch_size * beam_width,
                  min(
                      (ite + 1) * local_batch_size * beam_width,
                      batch_size * beam_width,
                  ),
              )
              if cache_indirection is not None:
                  bidx = range(
                      ite * local_batch_size,
                      min((ite + 1) * local_batch_size, batch_size),
                  )
                  src_cache_indirection = cache_indirection[
                      src_indir_idx, bidx, ...
                  ]
                  tgt_cache_indirection = cache_indirection[
                      tgt_indir_idx, bidx, ...
                  ]
  
              if step == max_input_length:
                  hidden_states = last_token_hidden_states[bbidx, ...]
              else:
                  if comm.is_pipeline_group_first():
                      input_embeds = self.word_embedding(
                          output_token_ids[step - 1, bbidx]
                      )
                      if self.position_encoding is not None:
                          position_ids = (step - 1) * torch.ones_like(
                              pad_lengths[bbidx]
                          )
                          input_embeds += self.position_encoding(
                              position_ids
                          )
                      if self.pre_decoder_layernorm is not None:
                          input_embeds = self.pre_decoder_layernorm(
                              input_embeds
                          )
                  else:
                      # Dummy input_imbeds
                      input_embeds = torch.empty(
                          size=(len(bbidx), self.decoder.hidden_size),
                          dtype=self.decoder.dtype,
                          device=device,
                      )
  
                  profiler.start("ft-decoder")
                  hidden_states = self.decoder.forward(
                      max_input_length=max_input_length,
                      step=step,
                      ite=ite,
                      input_embeds=input_embeds,
                      sequence_lengths=sequence_lengths[bbidx],
                      key_cache=k_cache,
                      value_cache=v_cache,
                      finished=finished[bbidx],
                      total_padding_tokens=pad_lengths[bbidx],
                      cache_indirection=src_cache_indirection,
                      masked_tokens=masked_tokens[bbidx, ...],
                  )
                  profiler.stop("ft-decoder")
  
              if comm.is_pipeline_group_last():
                  if self.post_decoder_layernorm is not None:
                      hidden_states = self.post_decoder_layernorm(
                          hidden_states
                      )
  
                  # We use logits of fp32 type to avoid overflow issue.
                  if self.use_fp32_to_compute_logit:
                      # The FT GPT op internally uses FP32 compute type
                      # for matrix multiplication.
                      # This will produce the same result with the
                      # end-to-end FT's GPT op.
                      logits = torch.nn.functional.linear(
                          hidden_states.float(), self.lm_head.weight
                      )
                  else:
                      logits = self.lm_head(hidden_states).float()
  
                  profiler.start("ft-decode")
                  should_stop = self.decode_op.forward(
                      logits.view(batch_size, beam_width, -1),
                      step,
                      max_input_length,
                      ite,
                      local_batch_size,
                      eos_token_ids,
                      top_k,
                      top_p,
                      temperature,
                      repetition_penalty,
                      presence_penalty,
                      min_length,
                      len_penalty,
                      beam_search_diversity_rate,
                      top_p_decay,
                      top_p_min,
                      top_p_reset_ids,
                      None,
                      input_lengths,
                      sequence_limit_lengths,
                      stop_words_list,
                      bad_words_list,
                      src_cache_indirection,
                      output_token_ids.view(-1, batch_size, beam_width),
                      finished,
                      sequence_lengths,
                      cum_log_probs,
                      output_log_probs,
                      parent_ids,
                      tgt_cache_indirection,
                  )
                  profiler.stop("ft-decode")
                  is_generation_done &= should_stop
  
          # Broadcast from the last pipeline node if needed.
          profiler.start("ft-bcast")
          tensors_to_bcast = [
              output_token_ids[step, ...],
              finished,
              sequence_lengths,
              is_generation_done,
          ]
          if beam_width > 1:
              tensors_to_bcast.append(tgt_cache_indirection)
          self.decode_op.broadcast_from_last_pipeline(tensors_to_bcast)
          profiler.stop("ft-bcast")
  
          if is_generation_done or finished.all():
              break
  
      # Transpose (L, batch, beam) -> (batch, beam, L)
      output_token_ids = output_token_ids.view(
          -1, batch_size, beam_width
      ).permute(1, 2, 0)
  
      # Increase sequence_length by 1 because the sequence length of time step t is t - 1. # noqa: E501
      sequence_lengths += 1
  
      # Outputs
      output_dict = dict(output_token_ids=output_token_ids)
      if return_output_length:
          output_dict["output_lengths"] = sequence_lengths
      if return_log_probs:
          output_dict["cum_log_probs"] = cum_log_probs
          output_dict["output_log_probs"] = output_log_probs
      return output_dict
  
imports: import math
  from abc import abstractmethod
  from pathlib import Path
  from typing import List, Literal, Optional, Union
  import os
  import numpy
  from .None import comm
  from .None import profiler
  from .gpt import GptInitModelParameters
  from nebullvm.optional_modules.torch import torch
label_desc: Use image node_id 14 for calling the Gpt obj's underlying member method code with example usage: obj.generate(input_token_ids, input_lengths, gen_length, eos_token_id, local_batch_size, beam_width, top_k, top_p, top_p_decay, top_p_min, top_p_reset_ids, temperature, repetition_penalty, presence_penalty, min_length, len_penalty, beam_search_diversity_rate, stop_words_list, bad_words_list, sequence_limit_lengths, random_seed, memory_length, return_output_length, return_log_probs) and returns: output_dict