class_docstr: Model class for prompt-tuning or p-tuning a pretrained Megatron model. 
  
  Prompt Tuning initalizes virtual prompt embeddings directly from a copy of
  certain token embeddings from the the pretrained 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__(self, cfg: DictConfig, trainer: Trainer):
      super().__init__(cfg, trainer)
  
      self.config: ModelParallelConfig = self.model_parallel_config
  
      self.load_frozen_model(cfg, trainer)
      self.prompt_encoder = None
      self.tokenizer = self.frozen_model.tokenizer
  
      if hasattr(self.frozen_model.cfg, "encoder") and hasattr(
          self.frozen_model.cfg, "decoder"
      ):
          self.hidden_size = (
              self.frozen_model.cfg.encoder.hidden_size
          )  # Encoder and decoder need to have the same hidden size and we check for this in the frozen enc-dec model.
          self.config.hidden_size = self.hidden_size
      else:
          self.hidden_size = self.frozen_model.cfg.hidden_size
          self.config.hidden_size = self.hidden_size
  
      self.existing_tasks = list(self.cfg.get("existing_tasks", []))
      self.new_tasks = list(self.cfg.get("new_tasks", []))
      self.virtual_prompt_style = VirtualPromptStyle(
          cfg.virtual_prompt_style
      )
  
      # 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,
          ]
      ):
          # TODO: Handle this when moving GPT prompt learning to the base class.
          self.word_embeddings = (
              self.frozen_model.enc_dec_model.encoder_embedding.word_embeddings
          )
  
      # 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
  
      # Prepare pseudo token ids for virtual/virtual prompt tokens
      self.pseudo_tokens = get_pseudo_tokens(self.max_virtual_tokens)
      if isinstance(self.tokenizer, SentencePieceTokenizer):
          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
      )
      self.decoder_seq_length = cfg.get("decoder_seq_length", 40)
  
      # 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
  
      self.enable_autocast = (
          True
          if (not self.megatron_amp_O2)
          and (self.autocast_dtype in [torch.float16, torch.bfloat16])
          else False
      )
  
      # 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 re
  from collections import OrderedDict
  from typing import Any, Optional
  import torch
  from omegaconf.dictconfig import DictConfig
  from omegaconf.omegaconf import open_dict
  from pytorch_lightning.trainer.trainer import Trainer
  from torch import Tensor
  from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer
  from nemo.collections.nlp.metrics.prompt_learning_metrics import AccuracyScore, BLEUScore, ROUGEScores
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common import PromptEncoder, PromptEncoderType, VirtualPromptPlaceholderToken, VirtualPromptSource, VirtualPromptStyle
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
  from nemo.collections.nlp.modules.common.transformer.text_generation import TextGeneration
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.utils import AppState, logging
label_desc: Use image node_id 1 to create a new MegatronBasePromptLearningModel object from inherited base classes: MegatronBaseModel, TextGeneration with example: obj = MegatronBasePromptLearningModel(cfg, trainer)