code: def test_wikitext2(self):
      batch_size = 7
      dataset_name = "wikitext"
      dataset_config_name = "wikitext-2-raw-v1"
      data_dir = Path(
          os.getenv("DATA_DIR", current_dir.parent.parent / "data")
      )
      cache_dir = data_dir / "wikitext-2" / "cache"
      max_length = 512
      dupe_factor = 2
      datamodule = MLMDataModule(
          dataset_name,
          tokenizer_name="bert-base-uncased",
          dataset_config_name=dataset_config_name,
          max_length=max_length,
          cache_dir=cache_dir,
          dupe_factor=dupe_factor,
          batch_size=batch_size,
          num_workers_preprocess=4,
      )
      datamodule.prepare_data()
      datamodule.setup(stage="fit")
      train_loader = datamodule.train_dataloader()
      val_loader = datamodule.val_dataloader()
      datamodule.setup(stage="test")
      test_loader = datamodule.test_dataloader()
      train_len = 69293
      val_len = 7222
      test_len = 8353
      assert len(train_loader) == div_up(train_len, batch_size)
      assert len(val_loader) == div_up(val_len, batch_size)
      assert len(test_loader) == div_up(test_len, batch_size)
      for loader in [train_loader, val_loader, test_loader]:
          batch = next(iter(loader))
          assert batch.keys() == {
              "input_ids",
              "labels",
              "attention_mask",
              "token_type_ids",
              "next_sentence_label",
          }
          seqlen = batch["input_ids"].shape[-1]
          assert batch["input_ids"].shape == (batch_size, seqlen)
          assert batch["input_ids"].dtype == torch.long
          assert batch["labels"].shape == (batch_size, seqlen)
          assert batch["labels"].dtype == torch.long
          assert batch["attention_mask"].shape == (batch_size, seqlen)
          assert batch["attention_mask"].dtype == torch.long
          assert batch["token_type_ids"].shape == (batch_size, seqlen)
          assert batch["token_type_ids"].dtype == torch.long
          assert batch["next_sentence_label"].shape == (batch_size,)
          assert batch["next_sentence_label"].dtype == torch.bool
  
file_path: zoo/tests/datamodules/test_masked_language_modeling.py
imports: import os
  from pathlib import Path
  import pytest
  import torch
  from src.datamodules.masked_language_modeling import MLMDataModule
label_desc: Use image node_id 1 for calling the TestMLMDataModule obj's underlying member method code with example usage: obj.test_wikitext2() without return types