code: def main(cfg) -> None:
      plugins = []
      if cfg.get("cluster_type", None) == "BCP":
          plugins.append(TorchElasticEnvironment())
      # trainer required for restoring model parallel models
      trainer = Trainer(
          plugins=plugins, strategy=NLPDDPStrategy(), **cfg.trainer
      )
  
      if (
          cfg.tensor_model_parallel_size < 0
          or cfg.pipeline_model_parallel_size < 0
          or cfg.get("pipeline_model_parallel_split_rank", -1) < 0
      ):
          model_config = MegatronNevaModel.restore_from(
              restore_path=cfg.neva_model_file,
              trainer=trainer,
              return_config=True,
          )
  
          with open_dict(cfg):
              cfg.tensor_model_parallel_size = model_config.get(
                  "tensor_model_parallel_size", 1
              )
              cfg.pipeline_model_parallel_size = model_config.get(
                  "pipeline_model_parallel_size", 1
              )
              cfg.pipeline_model_parallel_split_rank = model_config.get(
                  "pipeline_model_parallel_split_rank", 0
              )
  
      assert (
          cfg.trainer.devices * cfg.trainer.num_nodes
          == cfg.tensor_model_parallel_size
          * cfg.pipeline_model_parallel_size
      ), "devices * num_nodes should equal tensor_model_parallel_size * pipeline_model_parallel_size"
  
      if cfg.neva_model_file:
          save_restore_connector = NLPSaveRestoreConnector()
          if os.path.isdir(cfg.neva_model_file):
              save_restore_connector.model_extracted_dir = (
                  cfg.neva_model_file
              )
  
          neva_cfg = MegatronNevaModel.restore_from(
              restore_path=cfg.neva_model_file,
              trainer=trainer,
              return_config=True,
              save_restore_connector=save_restore_connector,
          )
          OmegaConf.set_struct(neva_cfg, True)
          with open_dict(neva_cfg):
              neva_cfg.sequence_parallel = False
              neva_cfg.activations_checkpoint_granularity = None
              neva_cfg.activations_checkpoint_method = None
              neva_cfg.precision = trainer.precision
              neva_cfg.mm_cfg.llm.from_pretrained = cfg.get(
                  "llm_model_file", None
              )
          #    neva_cfg.mm_cfg.vision_encoder.from_pretrained = None
  
          model = MegatronNevaModel.restore_from(
              restore_path=cfg.neva_model_file,
              trainer=trainer,
              override_config_path=neva_cfg,
              save_restore_connector=save_restore_connector,
          )
          if neva_cfg.get("peft") is not None:
              peft_cfg_cls = PEFT_CONFIG_MAP[neva_cfg.peft.peft_scheme]
              if peft_cfg_cls is not None:
                  model.load_adapters(
                      cfg.neva_model_file, peft_cfg_cls(neva_cfg)
                  )
  
      elif cfg.checkpoint_dir:
          app_state = AppState()
          if (
              cfg.tensor_model_parallel_size > 1
              or cfg.pipeline_model_parallel_size > 1
          ):
              app_state.model_parallel_size = (
                  cfg.tensor_model_parallel_size
                  * cfg.pipeline_model_parallel_size
              )
              app_state.tensor_model_parallel_size = (
                  cfg.tensor_model_parallel_size
              )
              app_state.pipeline_model_parallel_size = (
                  cfg.pipeline_model_parallel_size
              )
              (
                  app_state.tensor_model_parallel_rank,
                  app_state.pipeline_model_parallel_rank,
                  app_state.model_parallel_size,
                  app_state.data_parallel_size,
                  app_state.pipeline_model_parallel_split_rank,
                  app_state.virtual_pipeline_model_parallel_rank,
              ) = fake_initialize_model_parallel(
                  world_size=app_state.model_parallel_size,
                  rank=trainer.global_rank,
                  tensor_model_parallel_size_=cfg.tensor_model_parallel_size,
                  pipeline_model_parallel_size_=cfg.pipeline_model_parallel_size,
                  pipeline_model_parallel_split_rank_=cfg.pipeline_model_parallel_split_rank,
              )
          checkpoint_path = inject_model_parallel_rank(
              os.path.join(cfg.checkpoint_dir, cfg.checkpoint_name)
          )
          # TODO: This wont work properly (We need to set model.llm.from_pretrained model.vision.from_pretrained to nul)
          model = MegatronNevaModel.load_from_checkpoint(
              checkpoint_path,
              hparams_file=cfg.hparams_file,
              trainer=trainer,
          )
      else:
          raise ValueError(
              "need at least a nemo file or checkpoint dir"
          )
  
      model.freeze()
  
      # Have to turn off activations_checkpoint_method for inference
      try:
          model.model.language_model.encoder.activations_checkpoint_method = (
              None
          )
      except AttributeError:
          pass
      try:
          model.model.module.language_model.encoder.activations_checkpoint_method = (
              None
          )
      except AttributeError:
          pass
  
      length_params: LengthParam = {
          "max_length": cfg.inference.tokens_to_generate,
          "min_length": cfg.inference.min_tokens_to_generate,
      }
  
      sampling_params: SamplingParam = {
          "use_greedy": cfg.inference.greedy,
          "temperature": cfg.inference.temperature,
          "top_k": cfg.inference.top_k,
          "top_p": cfg.inference.top_p,
          "repetition_penalty": cfg.inference.repetition_penalty,
          "add_BOS": cfg.inference.add_BOS,
          "all_probs": cfg.inference.all_probs,
          "compute_logprob": cfg.inference.compute_logprob,
          "end_strings": cfg.inference.end_strings,
      }
  
      with open(cfg.prompt_file, "r") as f:
          lines = f.readlines()
  
      final_prompts = []
      for line in lines:
          prompt_dict = json.loads(line)
          final_prompts.append(prompt_dict)
  
      responses = model.generate(
          input_prompts=final_prompts,
          length_params=length_params,
          sampling_params=sampling_params,
          inference_config=cfg,
      )
  
      # =================== Start Quantization ====================
      if HAVE_AMMO and cfg.quantization.enable == True:
          print(
              f"Using quantization algorithm: {cfg.quantization.algorithm}"
          )
          if cfg.quantization.algorithm == "int8_sq":
              atq_config = atq.INT8_SMOOTHQUANT_CFG
          elif cfg.quantization.algorithm == "fp8":
              atq_config = atq.FP8_DEFAULT_CFG
          elif cfg.quantization.algorithm == "awq":
              atq_config = atq.INT4_AWQ_CFG
          else:
              raise ValueError(
                  f"Unsupported quantization algorithm: {cfg.quantization.algorithm}"
              )
  
          def forward_loop():
              model.generate(
                  input_prompts=final_prompts,
                  length_params=length_params,
                  sampling_params=sampling_params,
                  inference_config=cfg,
              )
  
          atq.quantize(model, atq_config, forward_loop)
  
          responses = model.generate(
              input_prompts=final_prompts,
              length_params=length_params,
              sampling_params=sampling_params,
              inference_config=cfg,
          )
      # ============== Quantization End =========================
  
      results = []
      for response, prompt in zip(responses, final_prompts):
          prompt["full_text"] = response["clean_text"]
          prompt["text"] = response["clean_response"]
          prompt["model_id"] = cfg.neva_model_file
          prompt["answer_id"] = 0
          prompt["metadata"] = {}
          results.append(prompt)
  
      with open(cfg.output_file, "w") as f:
          for result in results:
              f.write(json.dumps(result) + "\n")
  
      """ 
      # Second method of running text generation, call trainer.predict
      ds = RequestDataSet(final_prompts)
      request_dl = DataLoader(dataset=ds, batch_size=1)
      config = OmegaConf.to_container(cfg.inference)
      model.set_inference_config(config)
      response = trainer.predict(model, request_dl)
  
      print("***************************")
      print(response)
      print("***************************")
      """
  
imports: import json
  import os
  import torch
  from omegaconf import OmegaConf, open_dict
  from pytorch_lightning.plugins.environments import TorchElasticEnvironment
  from pytorch_lightning.trainer.trainer import Trainer
  from torch.utils.data import Dataset
  from nemo.collections.multimodal.models.multimodal_llm.neva.neva_model import MegatronNevaModel
  from nemo.collections.nlp.modules.common.megatron.megatron_init import fake_initialize_model_parallel
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, SamplingParam
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy, NLPSaveRestoreConnector
  from nemo.collections.nlp.parts.peft_config import PEFT_CONFIG_MAP
  from nemo.core.config import hydra_runner
  from nemo.utils.app_state import AppState
  from nemo.utils.model_utils import inject_model_parallel_rank
label_desc: Use image node_id 1 for calling a global function with example usage: main(cfg) without return types