code: def main(
      ckpt_dir: str,
      interactive: bool = False,
      add_instruction_tags: bool = True,
  ):
      if "WORLD_SIZE" in os.environ:
          mp_size = int(os.environ["WORLD_SIZE"])
          local_rank = int(os.environ["LOCAL_RANK"])
      else:
          mp_size = 1
          local_rank = 0
  
      device = mp_utils.initialize(mp_size, local_rank)
  
      g = FastGen.build(ckpt_dir, GenArgs(), device)
  
      for prompts in get_prompts(interactive):
          if add_instruction_tags:
              prompts = [f"[INST]{prompt}[/INST]" for prompt in prompts]
  
          tokens = [g.tokenizer.encode(x) for x in prompts]
          stats, out_tokens = g.generate_all(
              tokens, use_cuda_graphs="NO_CUDA_GRAPHS" not in os.environ
          )
  
          if mp_utils.get_rank() == 0:
              for i, prompt in enumerate(prompts):
                  print(f"> {prompt}")
                  answer = g.tokenizer.decode(out_tokens[i])
                  print(answer)
                  print("---------------")
  
              for phase_stats in stats.phases:
                  print(phase_stats.show())
  
imports: import json
  import os
  import readline
  import sys
  import time
  from dataclasses import dataclass
  from pathlib import Path
  from typing import Iterable, Optional, Tuple, Union
  import fire
  import model
  import mp_utils
  import sample_utils
  import torch
  from stats import Stats
  from tokenizer import Tokenizer
  from xformers.ops.fmha.attn_bias import BlockDiagonalCausalWithOffsetPaddedKeysMask
label_desc: Use image node_id 2 for calling a global function with example usage: main(ckpt_dir, interactive, add_instruction_tags) without return types