code: def process_sentence_chunks(
      ds: MMapRetrievalIndexedDataset,
      tokenizer,
      chunk_size: int,
      warm_up_size: int,
      percent: float,
      stage: Union[int, None],
      workers: int,
      shard_id: int,
      total_shards: int,
  ):
      """
      This function takes chunked tokens from the retrieval dataset and map it back to text.
      In stage 0, it only loads the first `warm_up_size` chunks that is used for building the Faiss index structure.
      In other stages, in addition to the warm_up_size chunks, it also sends the chunked text and add their embeddings into the index.
      In stage 1, it divides the total work into `total_shards`, and process only at the `shard_id`. If the stage is None, it process all the chunks.
      """
      total_chunks = ds.chunks
      num_docs = len(ds._index.sizes)
      assert len(ds._index.sizes) == len(ds._index._chunk_id_start)
      if percent < 1.0:
          use_num_docs = int(num_docs * percent)
          logging.info(
              f"Use {use_num_docs} out of {num_docs} docs to build index"
          )
          total_chunks = ds._index._chunk_id_start[
              min(use_num_docs, num_docs - 1)
          ]
      logging.info(f"{total_chunks} chunks are used to build the index")
      start = 0
      if stage is None or stage == 0:
          beg = time.time()
          # only prepare the warmup batch for stage None and stage 0
          assert warm_up_size < total_chunks
          warm_chunk_ids = np.random.randint(
              0, total_chunks, warm_up_size
          )
          warm_up_slices = []
          for warm_up_id in warm_chunk_ids:
              warm_up_slices.append(
                  ds.get_chunk(warm_up_id, force_no_cont_ids=True)
              )
          with Pool(workers) as p:
              sentences = p.map(tokenizer.ids_to_text, warm_up_slices)
          end = time.time()
          logging.info(
              f"token-to-text {total_chunks} chunks takes {end-beg}"
          )
          queue.put((sentences, None))
          if stage == 0:
              # first the task for stage 0
              queue.put((None, None))
              return
      elif stage == 1:
          shard_size = total_chunks // total_shards
          splits = list(range(0, total_chunks, shard_size))
          if shard_id < total_shards - 1:
              start = splits[shard_id]
              total_chunks = splits[shard_id + 1]
          elif shard_id == total_shards - 1:
              start = splits[shard_id]
              total_chunks = total_chunks
          else:
              raise ValueError(f"{shard_id} bigger than {total_shards}")
          logging.info(
              f"shard_id {shard_id}, create index from chunk {start} to {total_chunks}"
          )
  
      threshold = 0.1
      with Pool(workers) as p:
          while start < total_chunks:
              if start / total_chunks > threshold:
                  logging.info(
                      f"sentence processing {start / total_chunks} is done"
                  )
                  threshold += 0.1
              slice_id = (start, min(start + chunk_size, total_chunks))
              beg = time.time()
              id_slices = ds.get_chunk(
                  slice(*slice_id), force_no_cont_ids=True
              )
              end = time.time()
              logging.info(f"load {chunk_size} chunks takes {end-beg}")
              start = min(start + chunk_size, total_chunks)
              sentences = p.map(tokenizer.ids_to_text, id_slices)
              end2 = time.time()
              logging.info(
                  f"tokenize {chunk_size} chunks takes {end2-end}"
              )
              queue.put((sentences, slice_id))
      queue.put((None, None))
  
imports: import argparse
  import multiprocessing
  import pathlib
  import sys
  import time
  from multiprocessing import Pool
  from typing import Union
  import faiss
  import numpy
  import torch
  from faiss.contrib.ondisk import merge_ondisk
  from sentence_transformers import SentenceTransformer
  from nemo.collections.nlp.data.language_modeling.megatron.indexed_retrieval_dataset import MMapRetrievalIndexedDataset
  from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling a global function with example usage: process_sentence_chunks(ds, tokenizer, chunk_size, warm_up_size, percent, stage, workers, shard_id, total_shards) without return types