code: def _sample(
      model: Actor,
      input_ids: torch.Tensor,
      max_length: int,
      early_stopping: bool = False,
      eos_token_id: Optional[int] = None,
      pad_token_id: Optional[int] = None,
      top_k: Optional[int] = None,
      top_p: Optional[float] = None,
      temperature: Optional[float] = None,
      prepare_inputs_fn: Optional[
          Callable[[torch.Tensor, Any], dict]
      ] = None,
      update_model_kwargs_fn: Optional[
          Callable[[dict, Any], dict]
      ] = None,
      **model_kwargs,
  ) -> torch.Tensor:
      if input_ids.size(1) >= max_length:
          return input_ids
  
      logits_processor = _prepare_logits_processor(
          top_k, top_p, temperature
      )
      unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
  
      for _ in range(input_ids.size(1), max_length):
          model_inputs = (
              prepare_inputs_fn(input_ids, **model_kwargs)
              if prepare_inputs_fn is not None
              else {"input_ids": input_ids}
          )
          outputs = model(**model_inputs)
  
          # NOTE: this is correct only in left padding mode
          next_token_logits = outputs["logits"][:, -1, :]
          next_token_logits = logits_processor(
              input_ids, next_token_logits
          )
          # sample
          probs = torch.softmax(
              next_token_logits, dim=-1, dtype=torch.float
          )
          next_tokens = torch.multinomial(probs, num_samples=1).squeeze(
              1
          )
  
          # finished sentences should have their next token be a padding token
          if eos_token_id is not None:
              assert (
                  pad_token_id is not None
              ), "If `eos_token_id` is defined, make sure that `pad_token_id` is defined."
              next_tokens = (
                  next_tokens * unfinished_sequences
                  + pad_token_id * (1 - unfinished_sequences)
              )
  
          # update generated ids, model inputs for next step
          input_ids = torch.cat(
              [input_ids, next_tokens[:, None]], dim=-1
          )
          if update_model_kwargs_fn is not None:
              model_kwargs = update_model_kwargs_fn(
                  outputs, model_kwargs
              )
  
          # if eos_token was found in one sentence, set sentence to finished
          if eos_token_id is not None:
              unfinished_sequences = unfinished_sequences.mul(
                  (next_tokens != eos_token_id).long()
              )
  
          # stop when each sentence is finished if early_stopping=True
          if early_stopping and _is_sequence_finished(
              unfinished_sequences
          ):
              break
  
      return input_ids
  
file_path: ColossalAI/applications/Chat/coati/models/generation.py
imports: from typing import Any, Callable, Optional
  import torch
  import torch.distributed
  from transformers import PreTrainedTokenizer
  from .base import Actor
label_desc: Use image node_id 3 for calling a global function with example usage: _sample(model, input_ids, max_length, early_stopping, eos_token_id, pad_token_id, top_k, top_p, temperature, prepare_inputs_fn, update_model_kwargs_fn) and returns: input_ids, input_ids