code: def test_dataset_for_text_tokens(tmpdir):
      seed_everything(42)
  
      block_size = 1024 + 1
      cache = Cache(
          input_dir=str(tmpdir),
          chunk_size=block_size * 11,
          item_loader=TokensLoader(block_size),
      )
      text_idxs_list = []
  
      counter = 0
      while True:
          text_ids = torch.randint(
              0, 1000, (np.random.randint(0, 1000),)
          ).to(torch.int)
          text_idxs_list.append(text_ids)
          chunk_filepath = cache._add_item(counter, text_ids)
          if chunk_filepath:
              break
          counter += 1
  
      cache.done()
      cache.merge()
  
      dataset = StreamingDataset(
          input_dir=str(tmpdir), item_loader=TokensLoader(block_size)
      )
  
      assert len(dataset) == 10
  
      cache_0 = dataset[0]
      cache_1 = dataset[1]
      cache_2 = dataset[2]
      cache_3 = dataset[3]
      assert len(cache_0) == block_size
      assert len(cache_1) == block_size
      assert not torch.equal(cache_0, cache[1])
      indices = torch.cat(text_idxs_list, dim=0)
      assert torch.equal(cache_0, indices[: len(cache_0)])
      assert torch.equal(
          cache_1, indices[len(cache_0) : len(cache_0) + len(cache_1)]
      )
  
      dataloader = DataLoader(
          StreamingDataset(
              input_dir=str(tmpdir),
              item_loader=TokensLoader(block_size),
          ),
          batch_size=2,
      )
  
      for batch_idx, batch in enumerate(dataloader):
          if batch_idx == 0:
              assert torch.equal(torch.stack([cache_0, cache_1]), batch)
          elif batch_idx == 1:
              assert torch.equal(torch.stack([cache_2, cache_3]), batch)
          else:
              break
  
imports: import os
  import sys
  from time import sleep
  from unittest import mock
  import numpy
  import pytest
  import torch
  from lightning import seed_everything
  from lightning.data.streaming import Cache, functions
  from lightning.data.streaming.dataloader import StreamingDataLoader
  from lightning.data.streaming.dataset import _INDEX_FILENAME, Dir, RemoteDir, StreamingDataset, _associate_chunks_to_workers, _replay_chunks_sampling, _replay_sampling, _should_replace_path, _try_create_cache_dir
  from lightning.data.streaming.item_loader import TokensLoader
  from lightning.data.streaming.shuffle import FullShuffle, NoShuffle
  from lightning.data.utilities.env import Environment, _DistributedEnv, _WorkerEnv
  from torch.utils.data import DataLoader
label_desc: Use image node_id 9 for calling a global function with example usage: test_dataset_for_text_tokens(tmpdir) without return types