class_docstr: Megatron Retrieval enhanced language model
class_docstr_tok: ["Megatron","Retrieval","enhanced","language","model"]
code: def build_train_valid_test_datasets(self):
      logging.info("Building RETRO datasets.")
      global_batch_size = (
          self.trainer.world_size
          * self.cfg.micro_batch_size
          // self.cfg.tensor_model_parallel_size
      )
      # Compute trianing micro-batch steps: total_global_batch_steps x grad_acumms_per_global_batch
      max_train_steps = (
          self.trainer.max_steps * self.trainer.accumulate_grad_batches
      )
      eval_iters = (
          max_train_steps // self.trainer.val_check_interval + 1
      ) * self.trainer.limit_val_batches
      test_iters = self.trainer.limit_test_batches
  
      train_valid_test_num_samples = [
          max_train_steps * global_batch_size,
          eval_iters * global_batch_size,
          test_iters * global_batch_size,
      ]
      if self.cfg.data.get("mock", False):
          (
              self._train_ds,
              self._validation_ds,
              self._test_ds,
          ) = build_mock_train_valid_test_datasets(
              cfg=self.cfg,
              trainer=self.trainer,
              splits_string=self.cfg.data.splits_string,
              tokenizer=self.tokenizer,
              mock_data_size=self.cfg.data.get("mock_data_size", 10000),
          )
      else:
          (
              self._train_ds,
              self._validation_ds,
              self._test_ds,
          ) = build_train_valid_test_datasets(
              cfg=self.cfg,
              trainer=self.trainer,
              data_prefix=self.cfg.data.data_prefix,
              data_impl=self.cfg.data.data_impl,
              splits_string=self.cfg.data.splits_string,
              train_valid_test_num_samples=train_valid_test_num_samples,
              seq_length=self.cfg.data.seq_length,
              seed=self.cfg.seed,
              skip_warmup=self.cfg.data.get("skip_warmup", True),
              tokenizer=self.tokenizer,
              retrieval_prefix=self.cfg.data.retrieval_prefix,
              knn_map_path=self.cfg.data.knn_index,
          )
      if self._train_ds is not None:
          logging.info(
              f"Length of train dataset: {len(self._train_ds)}"
          )
      if self._validation_ds is not None:
          logging.info(
              f"Length of val dataset: {len(self._validation_ds)}"
          )
      if self._test_ds is not None:
          logging.info(f"Length of test dataset: {len(self._test_ds)}")
      logging.info(f"Finished building RETRO datasets.")
      return self._train_ds, self._validation_ds, self._test_ds
  
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 10 for calling the MegatronRetrievalModel obj's underlying member method code with example usage: obj.build_train_valid_test_datasets() and returns: self,  self,  self