code: def create_masked_lm_predictions(
      tokens,
      vocab_id_list,
      vocab_id_to_token_dict,
      masked_lm_prob,
      cls_id,
      sep_id,
      mask_id,
      max_predictions_per_seq,
      np_rng,
      max_ngram_size=3,
      mean_ngram_size=None,
      whole_word_masking=True,
      favor_long_ngrams=False,
      permutation=False,
      geometric_dist=False,
      masking_style="bert",
      tokenizer_type="wordpiece",
      skip_masking_id=None,
  ):
      """Creates the predictions for the masked LM objective.
      Note: Tokens here are vocab ids and not text tokens."""
      if not geometric_dist and mean_ngram_size is not None:
          raise ValueError(
              f"Mean ngram size is only supported for geometric distribution."
          )
  
      cand_indexes = []
      # Note(mingdachen): We create a list for recording if the piece is
      # the starting piece of current token, where 1 means true, so that
      # on-the-fly whole word masking is possible.
      token_boundary = [0] * len(tokens)
      skip_mask_idx = (
          None  # Store the index of token that cannot be masked.
      )
      for i, token in enumerate(tokens):
          if token == skip_masking_id:
              skip_mask_idx = i
          if token == cls_id or token == sep_id:
              token_boundary[i] = 1
              continue
          # Whole Word Masking means that if we mask all of the wordpieces
          # corresponding to an original word.
          #
          # Note that Whole Word Masking does *not* change the training code
          # at all -- we still predict each WordPiece independently, softmaxed
          # over the entire vocabulary.
          if (
              whole_word_masking
              and len(cand_indexes) >= 1
              and not is_start_piece(vocab_id_to_token_dict[token])
          ):
              cand_indexes[-1].append(i)
          else:
              cand_indexes.append([i])
              if is_start_piece(vocab_id_to_token_dict[token]):
                  token_boundary[i] = 1
  
      output_tokens = list(tokens)
  
      masked_lm_positions = []
      masked_lm_labels = []
  
      if masked_lm_prob == 0:
          return (
              output_tokens,
              masked_lm_positions,
              masked_lm_labels,
              token_boundary,
          )
  
      num_to_predict = min(
          max_predictions_per_seq,
          max(1, int(round(len(tokens) * masked_lm_prob))),
      )
      if masking_style != "bert":
          num_to_predict = max(1, num_to_predict)
      if num_to_predict < 1:
          logging.warning(
              f"Number of tokens is : {len(tokens)} and mask_probability is {masked_lm_prob}. None of the tokens will be masked"
          )
  
      ngrams = np.arange(1, max_ngram_size + 1, dtype=np.int64)
      if not geometric_dist:
          # Note(mingdachen):
          # By default, we set the probilities to favor shorter ngram sequences.
          pvals = 1.0 / np.arange(1, max_ngram_size + 1)
          pvals /= pvals.sum(keepdims=True)
          if favor_long_ngrams:
              pvals = pvals[::-1]
  
      ngram_indexes = []
      for idx in range(len(cand_indexes)):
          ngram_index = {}
          for n in ngrams:
              # Skip this ngram if it contains the index of token that should not be masked.
              # TODO: (sandeepsub) Generalize this to be a list of tokens that cannot be masked.
              if (
                  skip_mask_idx is not None
                  and skip_mask_idx >= idx
                  and skip_mask_idx <= idx + n
              ):
                  continue
              ngram_index[n] = cand_indexes[idx : idx + n]
          ngram_indexes.append(ngram_index)
  
      np_rng.shuffle(ngram_indexes)
  
      (masked_lms, masked_spans) = ([], [])
      covered_indexes = set()
      for cand_index_set in ngram_indexes:
          if len(masked_lms) >= num_to_predict:
              break
          if not cand_index_set:
              continue
          # Note(mingdachen):
          # Skip current piece if they are covered in lm masking or previous ngrams.
          for index_set in cand_index_set[1]:
              for index in index_set:
                  if index in covered_indexes:
                      continue
  
          if not geometric_dist:
              # Not all ngrams are available because of skip_masking_id that prevents a certain ID from being masked.
              available_ngrams = list(cand_index_set.keys())
              # n - 1 because pvals is 0-indexed and available ngrams are 1-indexed.
              pvals_current = np.array(
                  [pvals[n - 1] for n in available_ngrams]
              )
              n = np_rng.choice(
                  available_ngrams,
                  p=pvals_current / pvals_current.sum(keepdims=True),
              )
          else:
              # Sampling "n" from the geometric distribution and clipping it to
              # the max_ngrams. Using p=0.2 default from the SpanBERT paper
              # https://arxiv.org/pdf/1907.10529.pdf (Sec 3.1)
  
              # The expectation of a geometric distribution is E[X] = 1 / p
              p = (
                  1 / mean_ngram_size
                  if mean_ngram_size is not None
                  else 0.2
              )
              n = min(np_rng.geometric(p), max_ngram_size)
              # n may not be in the candidate index set because of skip_masking_id.
              # we try to find the nearest one in the candidate index set.
              if n not in cand_index_set:
                  n = _truncate_to_nearest(cand_index_set, n)
          index_set = sum(cand_index_set[n], [])
          n -= 1
          # Note(mingdachen):
          # Repeatedly looking for a candidate that does not exceed the
          # maximum number of predictions by trying shorter ngrams.
          while len(masked_lms) + len(index_set) > num_to_predict:
              if n == 0:
                  break
              if n - 1 in cand_index_set:
                  index_set = sum(cand_index_set[n - 1], [])
              n -= 1
          # If adding a whole-word mask would exceed the maximum number of
          # predictions, then just skip this candidate.
          if len(masked_lms) + len(index_set) > num_to_predict:
              continue
          is_any_index_covered = False
          for index in index_set:
              if index in covered_indexes:
                  is_any_index_covered = True
                  break
          if is_any_index_covered:
              continue
          for index in index_set:
              covered_indexes.add(index)
              masked_token = None
              if masking_style == "bert":
                  # 80% of the time, replace with [MASK]
                  if np_rng.random() < 0.8:
                      masked_token = mask_id
                  else:
                      # 10% of the time, keep original
                      if np_rng.random() < 0.5:
                          masked_token = tokens[index]
                      # 10% of the time, replace with random word
                      else:
                          masked_token = vocab_id_list[
                              np_rng.randint(0, len(vocab_id_list))
                          ]
              elif masking_style == "t5":
                  masked_token = mask_id
              elif masking_style == "bart":
                  masked_token = mask_id
              else:
                  raise ValueError("invalid value of masking style")
  
              output_tokens[index] = masked_token
              masked_lms.append(
                  MaskedLmInstance(index=index, label=tokens[index])
              )
  
          masked_spans.append(
              MaskedLmInstance(
                  index=index_set,
                  label=[tokens[index] for index in index_set],
              )
          )
  
      assert len(masked_lms) <= num_to_predict
      np_rng.shuffle(ngram_indexes)
  
      select_indexes = set()
      if permutation:
          if skip_masking_id is not None:
              raise ValueError(
                  f"permutation=True is not supported when skip_masking_id is not None."
              )
          for cand_index_set in ngram_indexes:
              if len(select_indexes) >= num_to_predict:
                  break
              if not cand_index_set:
                  continue
              # Note(mingdachen):
              # Skip current piece if they are covered in lm masking or previous ngrams.
              for index_set in cand_index_set[0]:
                  for index in index_set:
                      if (
                          index in covered_indexes
                          or index in select_indexes
                      ):
                          continue
  
              n = np.random.choice(
                  ngrams[: len(cand_index_set)],
                  p=pvals[: len(cand_index_set)]
                  / pvals[: len(cand_index_set)].sum(keepdims=True),
              )
              index_set = sum(cand_index_set[n - 1], [])
              n -= 1
  
              while (
                  len(select_indexes) + len(index_set) > num_to_predict
              ):
                  if n == 0:
                      break
                  index_set = sum(cand_index_set[n - 1], [])
                  n -= 1
              # If adding a whole-word mask would exceed the maximum number of
              # predictions, then just skip this candidate.
              if len(select_indexes) + len(index_set) > num_to_predict:
                  continue
              is_any_index_covered = False
              for index in index_set:
                  if (
                      index in covered_indexes
                      or index in select_indexes
                  ):
                      is_any_index_covered = True
                      break
              if is_any_index_covered:
                  continue
              for index in index_set:
                  select_indexes.add(index)
          assert len(select_indexes) <= num_to_predict
  
          select_indexes = sorted(select_indexes)
          permute_indexes = list(select_indexes)
          np_rng.shuffle(permute_indexes)
          orig_token = list(output_tokens)
  
          for src_i, tgt_i in zip(select_indexes, permute_indexes):
              output_tokens[src_i] = orig_token[tgt_i]
              masked_lms.append(
                  MaskedLmInstance(index=src_i, label=orig_token[src_i])
              )
  
      masked_lms = sorted(masked_lms, key=lambda x: x.index)
      # Sort the spans by the index of the first span
      masked_spans = sorted(masked_spans, key=lambda x: x.index[0])
  
      for p in masked_lms:
          masked_lm_positions.append(p.index)
          masked_lm_labels.append(p.label)
      return (
          output_tokens,
          masked_lm_positions,
          masked_lm_labels,
          token_boundary,
          masked_spans,
      )
  
imports: import collections
  import os
  import subprocess
  import time
  from typing import Any
  import numpy
  import torch
  from omegaconf import OmegaConf
  from omegaconf.dictconfig import DictConfig
  from nemo.collections.nlp.data.language_modeling.megatron.base_dataset_utils import get_datasets_weights_and_num_samples, get_train_valid_test_split_
  from nemo.collections.nlp.data.language_modeling.megatron.blendable_dataset import BlendableDataset
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_dataset import deallocate_indexed_dataset_memory
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_dataset import make_dataset
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_dataset import make_indexed_dataset_compatibility
  from nemo.collections.nlp.data.language_modeling.megatron.length_distribution_type import LengthDistribution
  from nemo.collections.nlp.data.language_modeling.megatron.lm_adapted_t5_dataset import T5LMAdaptedDataset
  from nemo.utils import logging
  from nemo.utils.get_rank import is_global_rank_zero
label_desc: Use image node_id 6 for calling a global function with example usage: create_masked_lm_predictions(tokens, vocab_id_list, vocab_id_to_token_dict, masked_lm_prob, cls_id, sep_id, mask_id, max_predictions_per_seq, np_rng, max_ngram_size, mean_ngram_size, whole_word_masking, favor_long_ngrams, permutation, geometric_dist, masking_style, tokenizer_type, skip_masking_id) and returns: output_tokens,  masked_lm_positions,  masked_lm_labels,  token_boundary,  masked_spans, output_tokens,  masked_lm_positions,  masked_lm_labels,  token_boundary