code: def create_data_iterator(
      mask_prob: float,
      random_replace_prob: float,
      unmask_replace_prob: float,
      batch_size: int,
      max_seq_length: int = 512,
      tokenizer: str = "roberta-base",
  ) -> InfiniteIterator:
      """Create the dataloader.
  
      Args:
          mask_prob (float):
              Fraction of tokens to mask
          random_replace_prob (float):
              Fraction of masked tokens to replace with random token
          unmask_replace_prob (float):
              Fraction of masked tokens to replace with the actual token
          batch_size (int):
              The batch size of the generated tensors
          max_seq_length (int, optional):
              The maximum sequence length for the MLM task. Defaults to 512.
          tokenizer (str, optional):
              The tokenizer to use. Defaults to "roberta-base".
  
      Returns:
          InfiniteIterator:
              The torch DataLoader, wrapped in an InfiniteIterator class, to
              be able to continuously generate samples
  
      """
      wikitext_dataset = datasets.load_dataset(
          "wikitext", "wikitext-2-v1", split="train"
      )
      wikitext_dataset = wikitext_dataset.filter(
          lambda record: record["text"] != ""
      ).map(lambda record: {"text": record["text"].rstrip("\n")})
      tokenizer = AutoTokenizer.from_pretrained(tokenizer)
      masking_function_partial = partial(
          masking_function,
          tokenizer=tokenizer,
          mask_prob=mask_prob,
          random_replace_prob=random_replace_prob,
          unmask_replace_prob=unmask_replace_prob,
          max_length=max_seq_length,
      )
      dataset = WikiTextMLMDataset(
          wikitext_dataset, masking_function_partial
      )
      collate_fn_partial = partial(
          collate_function, pad_token_id=tokenizer.pad_token_id
      )
      dataloader = DataLoader(
          dataset,
          batch_size=batch_size,
          shuffle=True,
          collate_fn=collate_fn_partial,
      )
  
      return InfiniteIterator(dataloader)
  
imports: import os
  import datetime
  import json
  import pathlib
  import re
  import string
  from functools import partial
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
  import random
  import datasets
  import fire
  import logging
  import loguru
  import numpy
  import pytz
  import sh
  import torch
  import torch.nn
  import deepspeed
  from torch.utils.data import DataLoader, Dataset
  from torch.utils.tensorboard import SummaryWriter
  from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
  from transformers.models.roberta import RobertaConfig, RobertaModel
  from transformers.models.roberta.modeling_roberta import RobertaLMHead, RobertaPreTrainedModel
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 5 for calling a global function with example usage: create_data_iterator(mask_prob, random_replace_prob, unmask_replace_prob, batch_size, max_seq_length, tokenizer) and returns: InfiniteIterator