class_docstr: Model class for prompt-tuning or p-tuning a pretrained Megatron GPT model. 
  
  Prompt Tuning initalizes virtual prompt embeddings directly from a copy of
  certain token embeddings from the the pretrained GPT model's vocabulary
  and directly tunes these embedding weights. The token embeddings used in 
  initalization are specified by the user in the config file. The model can 
  be prompt-tuned for multiple tasks at once. virtual prompts are stored in a 
  prompt table and can be added or deleted without disrupting virtual prompts 
  for other tasks. 
  
  P-tuning initializes an LSTM encoder model that generates virtual prompt
  embeddings for every task. Each task shares the same encoder. After ptuning
  is compelete, the learned virtual prompts can be saved to the prompt table
  using add_ptuned_prompts_to_prompt_table(). Thus, if a user wants to add a 
  new virtual prompt via p-tuning, they do not need to retrain on all previous 
  tasks. This gives p-tuning the same task flexiblity as prompt-tuning.
class_docstr_tok: ["Model","class","for","prompt-tuning","or","p-tuning","a","pretrained","Megatron","GPT","model",".","Prompt","Tuning","initalizes","virtual","prompt","embeddings","directly","from","a","copy","of","certain","token","embeddings","from","the","the","pretrained","GPT","model","'s","vocabulary","and","directly","tunes","these","embedding","weights",".","The","token","embeddings","used","in","initalization","are","specified","by","the","user","in","the","config","file",".","The","model","can","be","prompt-tuned","for","multiple","tasks","at","once",".","virtual","prompts","are","stored","in","a","prompt","table","and","can","be","added","or","deleted","without","disrupting","virtual","prompts","for","other","tasks",".","P-tuning","initializes","an","LSTM","encoder","model","that","generates","virtual","prompt","embeddings","for","every","task",".","Each","task","shares","the","same","encoder",".","After","ptuning","is","compelete",",","the","learned","virtual","prompts","can","be","saved","to","the","prompt","table","using","add_ptuned_prompts_to_prompt_table","(",")",".","Thus",",","if","a","user","wants","to","add","a","new","virtual","prompt","via","p-tuning",",","they","do","not","need","to","retrain","on","all","previous","tasks",".","This","gives","p-tuning","the","same","task","flexiblity","as","prompt-tuning","."]
code: def init_model(self, cfg: DictConfig, trainer: Trainer):
      self.cfg = cfg
      self.config: ModelParallelConfig = self.model_parallel_config
      save_restore_connector = NLPSaveRestoreConnector()
      if os.path.isdir(cfg.get("language_model_path")):
          save_restore_connector.model_extracted_dir = cfg.get(
              "language_model_path"
          )
      frozen_model_cfg = MegatronGPTModel.restore_from(
          cfg.get("language_model_path"),
          trainer=trainer,
          return_config=True,
          save_restore_connector=save_restore_connector,
      )
  
      # set hidden size in the model parallel config for pipeline parallel schedules
      setattr(self.config, "hidden_size", frozen_model_cfg.hidden_size)
  
      # Need to overwrite some params in frozen model's config before restoring
      with open_dict(frozen_model_cfg):
          frozen_model_cfg.megatron_amp_O2 = False
          frozen_model_cfg.optim.name = "fused_adam"
          frozen_model_cfg.micro_batch_size = self.cfg.micro_batch_size
          frozen_model_cfg.global_batch_size = (
              self.cfg.global_batch_size
          )
          frozen_model_cfg.precision = trainer.precision
          frozen_model_cfg.sequence_parallel = self.cfg.get(
              "sequence_parallel", False
          )
          frozen_model_cfg.activations_checkpoint_granularity = (
              self.cfg.get("activations_checkpoint_granularity", None)
          )
          frozen_model_cfg.activations_checkpoint_num_layers = (
              self.cfg.get("activations_checkpoint_num_layers", None)
          )
          frozen_model_cfg.activations_checkpoint_method = self.cfg.get(
              "activations_checkpoint_method", None
          )
  
      if cfg.get("language_model_path", None):
          self.frozen_model = MegatronGPTModel.restore_from(
              cfg.get("language_model_path"),
              trainer=trainer,
              save_restore_connector=save_restore_connector,
              override_config_path=frozen_model_cfg,
          ).to(dtype=self.autocast_dtype)
  
      self.megatron_amp_O2 = self.cfg.get("megatron_amp_O2", False)
      self.pipeline_parallel = (
          self.cfg.get("pipeline_model_parallel_size", 1) > 1
      )
      self.tokenizer = self.frozen_model.tokenizer
      self.hidden_size = self.frozen_model.cfg.hidden_size
      self.existing_tasks = list(self.cfg.get("existing_tasks", []))
      self.new_tasks = list(self.cfg.get("new_tasks", []))
      with open_dict(self.cfg):
          self.cfg.existing_tasks = (
              self.existing_tasks + self.new_tasks
          )  # TODO: for backward compatibility (@adithyare) in general these tasks lists should be deprecated
  
      self.virtual_prompt_style = VirtualPromptStyle(
          cfg.virtual_prompt_style
      )
      self.model_type = ModelType.encoder_or_decoder
  
      self.enable_autocast = (
          True
          if (not self.megatron_amp_O2)
          and (self.autocast_dtype in [torch.float16, torch.bfloat16])
          else False
      )
  
      if self.pipeline_parallel:
          assert (
              self.cfg.optim.sched.get("min_lr", 0.0) == 0.0
          ), "Minimum lr must be 0.0 when pipeline parallel size is > 1"
  
      # Load templates for assigning virtual prompt token positions
      self.load_task_templates(self.cfg.task_templates)
  
      if (
          self.first_stage_of_pipeline()
          and self.virtual_prompt_style
          in [
              VirtualPromptStyle.P_TUNING,
          ]
      ):
          if self.frozen_model.mcore_gpt:
              self.word_embeddings = (
                  self.frozen_model.model.embedding.word_embeddings
              )
          else:
              self.word_embeddings = (
                  self.frozen_model.model.language_model.embedding.word_embeddings
              )
  
      self.padded_vocab_size = self.frozen_model.padded_vocab_size
  
      # Prepare pseudo token ids for virtual/virtual prompt tokens
      self.pseudo_tokens = get_pseudo_tokens(self.max_virtual_tokens)
      if isinstance(self.tokenizer, SentencePieceTokenizer):
          if not self.tokenizer.legacy:
              if self.tokenizer.pad_id != -1:
                  self.tokenizer.pad_token = (
                      self.tokenizer.ids_to_tokens(
                          [self.tokenizer.pad_id]
                      )[0]
                  )
              else:
                  self.tokenizer.pad_token = (
                      self.tokenizer.ids_to_tokens(
                          [self.tokenizer.eos_id]
                      )[0]
                  )
              self.tokenizer.bos_token = self.tokenizer.ids_to_tokens(
                  [self.tokenizer.bos_id]
              )[0]
              self.tokenizer.eos_token = self.tokenizer.ids_to_tokens(
                  [self.tokenizer.eos_id]
              )[0]
              self.tokenizer.legacy = True
          self.tokenizer.add_special_tokens(self.pseudo_tokens)
      else:
          self.tokenizer.add_special_tokens(
              {"additional_special_tokens": self.pseudo_tokens}
          )
      self.pseudo_token_ids = self.tokenizer.tokens_to_ids(
          self.pseudo_tokens
      )
      self.pseudo_token_ids_start = (
          self.pseudo_token_ids[0] if self.pseudo_token_ids else None
      )
      self.pad_token_id = (
          self.tokenizer.pad_id
          if self.tokenizer.pad_id is not None
          else self.tokenizer.unk_id
      )
  
      # P-Tuning uses an LSTM Encoder to produce virtual token embeddings
      if self.virtual_prompt_style == VirtualPromptStyle.P_TUNING:
          self.virtual_prompt_source = (
              VirtualPromptSource.PROMPT_ENCODER
          )
      elif self.virtual_prompt_style == VirtualPromptStyle.NO_PROMPT:
          self.virtual_prompt_source = VirtualPromptSource.NO_PROMPT
      else:
          raise ValueError(
              f"\nvirtual prompt style '{cfg.virtual_prompt_style}.'"
          )
  
      self._reduced_loss_buffer = []
      self._inference_config = None
  
      # make sure the default pytorch lightning gradient clipping in the basemodel
      self.grad_clip_pl_default = True
      self.lowest_val_loss = None
      self.prompt_encoder = None
  
      # default inference related params -> for evaluation metrics
      if hasattr(self.cfg, "inference") and self.cfg.get(
          "report_validation_metric", False
      ):
          self.length_params: LengthParam = {
              "max_length": self.cfg.inference.get(
                  "tokens_to_generate", 30
              ),
              "min_length": self.cfg.inference.get(
                  "min_tokens_to_generate", 0
              ),
          }
  
          self.sampling_params: SamplingParam = {
              "use_greedy": self.cfg.inference.get("greedy", False),
              "temperature": self.cfg.inference.get("temperature", 1.0),
              "top_k": self.cfg.inference.get("tok_k", 0),
              "top_p": self.cfg.inference.get("top_p", 0.9),
              "repetition_penalty": self.cfg.inference.get(
                  "repetition_penalty", 1.2
              ),
              "add_BOS": True,
              "all_probs": False,
              "compute_logprob": False,
              "end_strings": self.cfg.inference.get(
                  "end_strings", ["<|endoftext|>"]
              ),
          }
      elif self.cfg.get(
          "report_validation_metric", False
      ) and not hasattr(self.cfg, "inference"):
          raise ValueError(
              "Must provide inference parameters for reporting validation metric!"
          )
  
      # define validation metric
      if self.cfg.get("report_validation_metric", False):
          validation_metric = self.cfg.get(
              "validation_metric", "accuracy"
          )
          if validation_metric == "accuracy":
              self.validation_metric = AccuracyScore()
          elif validation_metric == "bleu":
              self.validation_metric = BLEUScore()
          elif validation_metric == "rouge":
              self.validation_metric = ROUGEScores()
  
imports: import itertools
  import os
  from functools import partial
  from typing import Any, List, Optional, Union
  import torch
  from omegaconf.dictconfig import DictConfig
  from omegaconf.omegaconf import open_dict
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
  from nemo.collections.nlp.data.language_modeling.megatron.gpt_prompt_learning_dataset import GPTPromptLearningDataset
  from nemo.collections.nlp.metrics.prompt_learning_metrics import AccuracyScore, BLEUScore, ROUGEScores
  from nemo.collections.nlp.models.language_modeling.megatron_base_prompt_learning_model import MegatronBasePromptLearningModel
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
  from nemo.collections.nlp.modules.common import VirtualPromptPlaceholderToken, VirtualPromptSource, VirtualPromptStyle
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_iterator_k_split
  from nemo.collections.nlp.modules.common.text_generation_utils import get_default_length_params, get_default_sampling_params, megatron_gpt_generate
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, SamplingParam
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler, NLPSaveRestoreConnector
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 2 for calling the MegatronGPTPromptLearningModel obj's underlying member method code with example usage: obj.init_model(cfg, trainer) without return types