code: def process_dataset(self):
      cache_dir = (
          None
          if self.cache_dir is None
          else self.cache_dir / self._cache_dir_name
      )
      if cache_dir is not None:
          if cache_dir.is_dir():
              return self._load_from_cache(cache_dir)
  
      raw_datasets = load_dataset(
          self.dataset_name, self.dataset_config_name
      )
      # https://github.com/stanford-crfm/mistral/blob/main/src/corpora/auto.py
      if "validation" not in raw_datasets:
          assert (
              "train" in raw_datasets
          ), "You must have train in raw_datasets to make a validation raw_datasets"
          raw_datasets = raw_datasets["train"].train_test_split(
              test_size=self.val_ratio,
              seed=self.val_split_seed,
              shuffle=True,  # Otherwise test will be at the end of the dataset
          )
          raw_datasets["validation"] = raw_datasets["test"]
  
      if (
          self.val_only
      ):  # Should only be used for evaluation, not for training
          raw_datasets["train"] = raw_datasets["validation"]
  
      # [2021-12-25] TD: Running the detokenizer on wikitext-103 makes ppl worse
      # (GPT2-small val ppl after 10 epochs ~22 -> ~25)
      # However, it's useful for zero-shot transfer from Openwebtext,
      # as after detokenization it's closer to Openwebtext's format.
      # https://github.com/stanford-crfm/mistral/issues/12
      if self.detokenize:
          if self.dataset_name in DATASET_TOKENIZATION_REGISTRY:
              detokenizer = DATASET_TOKENIZATION_REGISTRY[
                  self.dataset_name
              ]
              raw_datasets = raw_datasets.map(
                  lambda example: {
                      "text": detokenizer(example["text"])
                  },
                  num_proc=max(self.num_workers, 1),
                  desc="Running detokenizer on dataset",
              )
  
      tokenizer = AutoTokenizer.from_pretrained(
          self.tokenizer_name, use_fast=True
      )
      # Preprocessing the datasets.
      # First we tokenize all the texts.
      column_names = raw_datasets["train"].column_names
      text_column_name = (
          "text" if "text" in column_names else column_names[0]
      )
      # [2021-12-25] TD: For wikitext, don't need to add the EOS since each example already ends
      # with '\n', and there are no other '\n' in the examples.
      # assert all([t.count('\n') == 1 for t in raw_datasets['train']['text'] if t])
      # Add EOS token to the end of the text if the text is not empty
      # https://github.com/stanford-crfm/mistral/issues/91
      # https://github.com/stanford-crfm/mistral/pull/98
      if self.add_eos:
          add_eos = (
              lambda seq: (seq + tokenizer.eos_token) if seq else seq
          )
          add_eos_batched = lambda seqs: [add_eos(seq) for seq in seqs]
          tokenize = lambda example: tokenizer(
              add_eos_batched(example[text_column_name])
          )
      else:
          tokenize = lambda example: tokenizer(
              example[text_column_name]
          )
      # tokenized_datasets = raw_datasets.map(
      #     tokenize,
      #     batched=True,
      #     num_proc=max(self.num_workers, 1),
      #     remove_columns=column_names,
      #     desc="Running tokenizer on dataset",
      # )
      dtype = (
          np.uint16 if tokenizer.vocab_size < 64 * 1024 else np.int32
      )
  
      def tokenize_concat(examples):
          # We just need 'input_ids', not 'attention_mask' (since it's all 1)
          input_ids = np.fromiter(
              chain(*tokenize(examples)["input_ids"]), dtype=dtype
          )
          # Need to return a list since we're doing batched processing
          return {"input_ids": [input_ids], "len": [len(input_ids)]}
  
      tokenized_datasets = raw_datasets.map(
          tokenize_concat,
          batched=True,
          num_proc=max(self.num_workers, 1),
          remove_columns=column_names,
          desc="Running tokenizer on dataset",
      )
  
      if self.use_shmem:
          # Concatenate all input_ids into an array in shared memory
          def write_ids_to_shm(example, shm_name, array_len):
              shm = SharedMemory(name=shm_name)
              shm_arr = np.ndarray(
                  (array_len,), dtype=dtype, buffer=shm.buf
              )
              start_idx = example["len_offset"] - len(
                  example["input_ids"]
              )
              shm_arr[start_idx : example["len_offset"]] = example[
                  "input_ids"
              ]
              shm.close()
  
          concat_ids = {}
          for name, ds in tokenized_datasets.items():
              tokenized_datasets[name] = ds.add_column(
                  "len_offset", np.cumsum(ds["len"])
              )
              array_len = tokenized_datasets[name][-1]["len_offset"]
              shm = SharedMemory(
                  create=True, size=array_len * np.dtype(dtype).itemsize
              )
              shm_name = shm.name
              tokenized_datasets[name].map(
                  write_ids_to_shm,
                  fn_kwargs={
                      "shm_name": shm_name,
                      "array_len": array_len,
                  },
                  batched=False,
                  num_proc=max(self.num_workers, 1),
                  desc="Concatenating examples",
              )
              shm_arr = np.ndarray(
                  (array_len,), dtype=dtype, buffer=shm.buf
              )
              # We need to keep a reference to the shared memory, otherwise it gets garbage-collected
              # when it goes out of scope, and that memory is gone.
              # https://github.com/numpy/numpy/issues/18294
              concat_ids[name] = SHMArray(shm_arr, shm=shm)
      else:
          # Use disk
          concat_ids = {}
          assert cache_dir is not None
          cache_dir.mkdir(parents=True, exist_ok=True)
  
          def write_ids_to_disk(example, filename):
              with open(filename, "r+b") as f:
                  mm = mmap.mmap(f.fileno(), 0)
                  start_idx = example["len_offset"] - len(
                      example["input_ids"]
                  )
                  array_len = len(example["input_ids"])
                  arr = np.ndarray(
                      (array_len,),
                      dtype=dtype,
                      buffer=mm,
                      offset=np.dtype(dtype).itemsize * start_idx,
                  )
                  arr[:] = example["input_ids"]
                  mm.flush()
  
          for name, ds in tokenized_datasets.items():
              tokenized_datasets[name] = ds.add_column(
                  "len_offset", np.cumsum(ds["len"])
              )
              array_len = tokenized_datasets[name][-1]["len_offset"]
              filename = cache_dir / f"{name}.bin"
              # Need to create the file with this specific size first
              # https://ostechnix.com/create-files-certain-size-linux/
              subprocess.run(
                  [
                      "truncate",
                      "-s",
                      str(array_len * np.dtype(dtype).itemsize),
                      str(filename),
                  ],
                  check=True,
              )
              tokenized_datasets[name].map(
                  write_ids_to_disk,
                  fn_kwargs={"filename": filename},
                  batched=False,
                  num_proc=max(self.num_workers, 1),
                  desc="Concatenating examples",
              )
              concat_ids[name] = np.memmap(
                  filename, dtype=dtype, mode="r", shape=(array_len,)
              )
  
      if cache_dir is not None:
          self._save_to_cache(concat_ids, tokenizer, cache_dir)
          if not self.use_shmem:
              for name in concat_ids:
                  Path(cache_dir / f"{name}.bin").unlink()
      return concat_ids, tokenizer
  
imports: from itertools import chain
  from pathlib import Path
  import pickle
  from typing import Any, List, Union
  import subprocess
  import mmap
  from multiprocessing.shared_memory import SharedMemory
  import numpy
  import torch
  from torch.utils.data.dataloader import DataLoader, Dataset
  from transformers import AutoTokenizer
  from datasets import load_dataset
  from pytorch_lightning import LightningDataModule
  from src.datamodules.datasets.lm_dataset import LMDataset
  from src.datamodules.fault_tolerant_sampler import RandomFaultTolerantSampler
  from src.datamodules.fault_tolerant_sampler import FaultTolerantDistributedSampler
  from src.datamodules.datasets.detokenizer import DATASET_TOKENIZATION_REGISTRY
  from src.utils.utils import get_logger
label_desc: Use image node_id 4 for calling the LMDataModule obj's underlying member method code with example usage: obj.process_dataset() and returns: concat_ids,  tokenizer, dict, self