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.
docstr: Replaces the virtual tokens in the input_ids with embeddings 
  calculated from either the 'prompt_table' or 'prompt_encoder'. 
  The virtual token placeholders have token_ids listed in
  `self.pseudo_token_ids`.
  
  params:
      input_ids: the input token ids
      taskname_ids: the NLP task tag token ids
  returns:
      the token embedding for the LM model.
docstr_tok: ["Replaces","the","virtual","tokens","in","the","input_ids","with","embeddings","calculated","from","either","the","'prompt_table","'","or","'prompt_encoder","'",".","The","virtual","token","placeholders","have","token_ids","listed","in","`","self.pseudo_token_ids","`",".","params",":","input_ids",":","the","input","token","ids","taskname_ids",":","the","NLP","task","tag","token","ids","returns",":","the","token","embedding","for","the","LM","model","."]
code: def embed_input(
      self,
      input_ids: Tensor,
      taskname_ids: Tensor,
      use_cached_reps: bool,
  ):
      """
      Replaces the virtual tokens in the input_ids with embeddings
      calculated from either the 'prompt_table' or 'prompt_encoder'.
      The virtual token placeholders have token_ids listed in
      `self.pseudo_token_ids`.
  
      params:
          input_ids: the input token ids
          taskname_ids: the NLP task tag token ids
      returns:
          the token embedding for the LM model.
      """
      # Replace virtual token ids with padding for forward pass through vocab embeddings
      discrete_token_ids = input_ids.clone()
      discrete_token_ids[
          (input_ids >= self.pseudo_token_ids_start)
      ] = self.pad_token_id
      discrete_token_embeds = self.word_embeddings(
          discrete_token_ids
      ).clone()
  
      # Find the indicies where virtual tokens should be inserted
      virtual_token_locations = input_ids >= self.pseudo_token_ids_start
  
      # If there are no virtual tokens, just return discrete token embeds
      if not virtual_token_locations.any():
          return discrete_token_embeds
  
      if (
          self.virtual_prompt_source
          == VirtualPromptSource.PROMPT_ENCODER
      ):
          # taskname_embeddings = self.word_embeddings(taskname_ids)
          batch_size, _ = taskname_ids.size()
          virtual_token_embeds = self.prompt_encoder(
              batch_size=batch_size, use_cached_reps=use_cached_reps
          )
      else:
          raise ValueError("invalid VirtualPromptSource.")
  
      # Create index template specifying where virtual token embeddings should be placed
      batch_size, _, embedding_size = discrete_token_embeds.shape
      virtual_token_index = virtual_token_locations.nonzero().reshape(
          (batch_size, -1, 2)
      )[:, :, 1][:, :, None]
      virtual_token_index = virtual_token_index.expand(
          batch_size, self.total_new_task_virtual_tokens, embedding_size
      )
  
      # Make sure discrete_token_embeds and virtual_token_embeds share the same dtype
      discrete_token_embeds = discrete_token_embeds.type(
          virtual_token_embeds.dtype
      )
  
      # Insert virtual token embeddings where they belong amoung the discrete token embeddings
      discrete_token_embeds.scatter_(
          1, virtual_token_index, virtual_token_embeds
      )
      input_embeds = discrete_token_embeds
  
      return input_embeds
  
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 8 for calling the MegatronBasePromptLearningModel obj's underlying member method code with example usage: obj.embed_input(input_ids, taskname_ids, use_cached_reps) and returns: input_embeds, discrete_token_embeds