class_docstr: Megatron Retrieval enhanced language model
class_docstr_tok: ["Megatron","Retrieval","enhanced","language","model"]
code: def model_provider_func(
      self, pre_process, post_process, add_encoder, add_decoder
  ):
      # TODO: create get_encoder_decoder_model()here for different losses (e..g, nll, vae, mim)
  
      model = MegatronRetrievalTokenLevelEncoderDecoderModule(
          config=self.model_parallel_config,
          vocab_size=self.padded_vocab_size,
          hidden_size=self.cfg.hidden_size,
          max_position_embeddings=self.cfg.max_position_embeddings,
          num_attention_heads=self.cfg.num_attention_heads,
          ffn_hidden_size=self.cfg.ffn_hidden_size,
          apply_query_key_layer_scaling=self.cfg.get(
              "apply_query_key_layer_scaling", True
          ),
          kv_channels=self.cfg.get("kv_channels", None),
          num_tokentypes=0,
          parallel_output=True,
          pre_process=pre_process,
          post_process=post_process,
          init_method_std=self.cfg.get("init_method_std", 0.02),
          fp16_cross_entropy=self.cfg.get(
              "fp16_lm_cross_entropy", False
          ),
          hidden_dropout=self.cfg.get("hidden_dropout", 0.1),
          attention_dropout=self.cfg.get("attention_dropout", 0.1),
          precision=self.cfg.get("precision", 16),
          fp32_residual_connection=self.cfg.get(
              "fp32_residual_connection", False
          ),
          activations_checkpoint_method=self.cfg.get(
              "activations_checkpoint_method", None
          ),
          activations_checkpoint_num_layers=self.cfg.get(
              "activations_checkpoint_num_layers", 1
          ),
          layernorm_epsilon=self.cfg.get("layernorm_epsilon", 1e-5),
          persist_layer_norm=self.cfg.get("persist_layer_norm", False),
          bias_gelu_fusion=self.cfg.get("bias_gelu_fusion", True),
          bias_dropout_add_fusion=self.cfg.get(
              "bias_dropout_add_fusion", True
          ),
          masked_softmax_fusion=self.cfg.get(
              "masked_softmax_fusion", True
          ),
          onnx_safe=self.cfg.get("onnx_safe", False),
          activation=self.cfg.get("activation", "gelu"),
          bias=self.cfg.get("bias", True),
          normalization=self.cfg.get("normalization", "layernorm"),
          headscale=self.cfg.get("headscale", False),
          transformer_block_type=self.cfg.get(
              "transformer_block_type", "pre_ln"
          ),
          add_encoder=add_encoder,
          add_decoder=add_decoder,
          chunk_size=self.cfg.get(
              "chunk_size", 64
          ),  # the chunk size used to retrive
          enc_num_layers=self.cfg.get(
              "enc_num_layers", 4
          ),  # total number of encoder layers
          dec_num_layers=self.cfg.get(
              "dec_num_layers", 6
          ),  # total number of decoder layers
          enc_cross_attention=self.cfg.get(
              "enc_cross_attention", [3]
          ),  # layer numbers for cross attention
          dec_cross_attention=self.cfg.get(
              "dec_cross_attention", [3, 5]
          ),  # layer numbers for chunked cross attention
          add_position_embedding=self.cfg.get(
              "add_position_embedding", False
          ),  # whether use the absolute postion encoding
          tokenizer=self.tokenizer,
          activations_checkpoint_granularity=self.cfg.get(
              "activations_checkpoint_granularity", None
          ),
          megatron_lm_compatible=self.cfg.get(
              "megatron_lm_compatible", False
          ),
          version=self.cfg.get("version", 1),
      )
      return model
  
imports: import math
  from typing import Any, List, Optional, Union
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.data.language_modeling.megatron.retro_dataset import build_mock_train_valid_test_datasets, build_train_valid_test_datasets
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.mup.init import normal_
  from nemo.collections.nlp.modules.common.megatron.mup.shape import set_base_shapes
  from nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder import MegatronRetrievalTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group, build_position_ids, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_strategy import model_inference_strategy_dispatcher
  from nemo.collections.nlp.modules.common.text_generation_utils import generate, get_computeprob_response, get_default_length_params, get_default_sampling_params, megatron_gpt_generate
  from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, OutputType, SamplingParam, TextGeneration
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.utils import AppState, logging
label_desc: Use image node_id 3 for calling the MegatronRetrievalModel obj's underlying member method code with example usage: obj.model_provider_func(pre_process, post_process, add_encoder, add_decoder) and returns: model