docstr: Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
docstr_tok: ["Takes","a","conditioning","sequence","(","prompt",")","as","input","and","continues","to","generate","as","many","tokens","as","requested","."]
code: def generate(
      self,
      prompt: torch.Tensor,
      max_new_tokens: int,
      *,
      callback=lambda x: x,
      **sampling_kwargs,
  ) -> torch.Tensor:
      """
      Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.
      """
      # create an empty tensor of the expected final shape and fill in the current tokens
      T = prompt.size(0)
      T_new = T + max_new_tokens
  
      max_seq_length = min(T_new, self.model.config.block_size)
      max_seq_length = (
          max_seq_length + self.speculate_k + 1
          if self.is_speculative
          else max_seq_length
      )
  
      dtype = prompt.dtype
      with torch.device(self.device):
          self.model.setup_caches(
              max_batch_size=1, max_seq_length=max_seq_length
          )
          if self.is_speculative and self.draft_model is not self.model:
              self.draft_model.setup_caches(
                  max_batch_size=1, max_seq_length=max_seq_length
              )
  
      # create an empty tensor of the expected final shape and fill in the current tokens
      empty = torch.empty(T_new, dtype=dtype, device=self.device)
      empty[:T] = prompt
      seq = empty
      input_pos = torch.arange(0, T, device=self.device)
  
      next_token = self.prefill(
          self.model, prompt.view(1, -1), input_pos, **sampling_kwargs
      )
      if self.is_speculative:
          self.prefill(
              self.draft_model,
              prompt.view(1, -1),
              input_pos,
              **sampling_kwargs,
          )
  
      period_id = self.tokenizer.encode(".")[0]
      text = self.tokenizer.decode([period_id] + next_token.tolist())[
          1:
      ]
      if self.stream:
          send_intermediate_predict_response(
              [text],
              self.context.request_ids,
              "Intermediate Prediction success",
              200,
              self.context,
          )
  
      seq[T] = next_token
  
      input_pos = torch.tensor([T], device=self.device, dtype=torch.int)
      accept_counts = [0] * (self.speculate_k + 1)
  
      if self.is_speculative:
          input_pos = (
              input_pos.item()
          )  # for speculative decoding easier to keep on host
          while input_pos < T_new - 1:
              cur_token = next_token.view(())
  
              next_tokens = self.speculative_decode(
                  self.model,
                  self.draft_model,
                  cur_token,
                  input_pos,
                  self.speculate_k,
                  **sampling_kwargs,
              )
  
              accept_counts[len(next_tokens) - 1] += 1
              num_added = min(T_new - input_pos - 1, len(next_tokens))
              seq[
                  input_pos + 1 : input_pos + num_added + 1
              ] = next_tokens[:num_added]
              for i in next_tokens[:num_added,]:
                  callback(i)
              input_pos = input_pos + num_added
              next_token = next_tokens[-1]
      else:
          generated_tokens, _ = self.decode_n_tokens(
              next_token.view(1, -1),
              input_pos,
              max_new_tokens - 1,
              callback=callback,
              **sampling_kwargs,
          )
          seq[T + 1 :] = torch.cat(generated_tokens)
  
      generate_stats = {"accept_counts": accept_counts}
  
      return seq, generate_stats
  
file_path: serve/examples/large_models/gpt_fast/handler.py
imports: import json
  import logging
  import time
  from pathlib import Path
  import torch
  from generate import Transformer, _load_model, decode_one_token, encode_tokens, logits_to_probs, maybe_init_dist, model_forward, multinomial_sample_one_no_sync, prefill
  from sentencepiece import SentencePieceProcessor
  from ts.handler_utils.timer import timed
  from ts.handler_utils.utils import send_intermediate_predict_response
  from ts.torch_handler.base_handler import BaseHandler
label_desc: Use image node_id 6 for calling the GptHandler obj's underlying member method code with example usage: obj.generate(prompt, max_new_tokens) and returns: seq,  generate_stats