code: def create_extreme_masked_lm_predictions(
      tokens,
      masked_lm_prob,
      mask_id,
      max_predictions_per_seq,
      np_rng,
      max_ngram_size=10,
      min_ngram_size=2,
      mean_ngram_size=5,
      span_length_distribution=LengthDistribution.uniform,
      skip_masking_id=None,
  ):
      """Creates the predictions for the extreme span-masking UL2 objective.
      Note: Tokens here are vocab ids and not text tokens."""
      output_tokens = list(tokens)
  
      masked_lm_positions = []
      masked_lm_labels = []
  
      num_to_predict = int(
          min(
              max_predictions_per_seq,
              max(1, int(round(len(tokens) * masked_lm_prob))),
          )
      )
      # If the number of tokens to predict is less than the min ngram size, clam it to max predictions.
      min_ngram_size = int(min(num_to_predict, min_ngram_size))
  
      ngrams = np.arange(
          min_ngram_size, max_ngram_size + 1, dtype=np.int64
      )
      if span_length_distribution == "uniform":
          pvals = np.array(
              [1.0 / (max_ngram_size - min_ngram_size + 1)]
              * (max_ngram_size - min_ngram_size + 1)
          )
  
      ngram_indexes = []
      if skip_masking_id is not None:
          skip_mask_idx = None
          for idx in range(len(tokens)):
              if tokens[idx] == skip_masking_id:
                  skip_mask_idx = idx
                  break
      else:
          skip_mask_idx = None
  
      cand_indexes = [[i] for i in range(len(tokens))]
      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[min_ngram_size]:
              for index in index_set:
                  if index in covered_indexes:
                      continue
  
          if span_length_distribution == LengthDistribution.uniform:
              available_ngrams = list(cand_index_set.keys())
              pvals_current = np.array(
                  [pvals[n] for n in available_ngrams]
              )
              n = np_rng.choice(
                  available_ngrams,
                  p=pvals_current / pvals_current.sum(keepdims=True),
              )
          elif span_length_distribution == LengthDistribution.geometric:
              # 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)
              n = int(np.clip(n, min_ngram_size, max_ngram_size))
          elif (
              span_length_distribution
              == LengthDistribution.truncated_normal
          ):
              # Sampling "n" from a truncated normal distribution.
              mu = (
                  mean_ngram_size
                  if mean_ngram_size is not None
                  else (max_ngram_size - min_ngram_size) // 2
              )
              n = int(
                  np.clip(
                      np_rng.normal(loc=mu, scale=np.sqrt(mu)),
                      min_ngram_size,
                      max_ngram_size,
                  )
              )
              if n not in cand_index_set:
                  n = _truncate_to_nearest(cand_index_set, n)
                  n = int(np.clip(n, min_ngram_size, max_ngram_size))
  
          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 < min_ngram_size:
                  break
              if n in cand_index_set:
                  index_set = sum(cand_index_set[n], [])
              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)
              output_tokens[index] = mask_id
              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)
  
      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,
          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 8 for calling a global function with example usage: create_extreme_masked_lm_predictions(tokens, masked_lm_prob, mask_id, max_predictions_per_seq, np_rng, max_ngram_size, min_ngram_size, mean_ngram_size, span_length_distribution, skip_masking_id) and returns: output_tokens,  masked_lm_positions,  masked_lm_labels,  masked_spans