code: def _train_sentencepiece(
      dataset: tf.data.Dataset,
      *,
      vocab_size: int,
      maxchars: int = int(1e7),
      model_path: str,
      model_type: str = "unigram",
      character_coverage: float = 1.0,
      data_keys=("inputs", "targets"),
  ):
      """Train SentencePiece tokenizer from subset of tf dataset.
  
      Args:
        dataset: tf.dataset
        vocab_size: int: size of vocab tokens to train.
        maxchars: int: number of characters to use for sentencepiece training.
        model_path: str: path of model file to save vocab model to.
        model_type: str: type of sentencepiece vocab to train.
        character_coverage: amount of characters covered by the model, good defaults
          are 0.9995 for languages with rich character set like Japanese or Chinese
          and 1.0 for other languages with small character set.
        data_keys: Tuple[str]: keys of dataset to use for training.
  
      Returns:
        path to the trained sentencepiece vocabulary model.
      """
      if model_path.startswith("gs://"):
          abs_model_path = model_path
      else:
          abs_model_path = os.path.abspath(
              os.path.expanduser(model_path)
          )
      fname, _ = _dump_chars_to_textfile(
          dataset, maxchars=maxchars, data_keys=data_keys
      )
      with tempfile.NamedTemporaryFile(
          delete=False, prefix="/tmp/sp_tmp"
      ) as model_fp:
          pass  # we just want a prefix'd tmp-filename
      argstr = " ".join(
          [
              f"--input={fname}",
              f"--vocab_size={vocab_size}",
              f"--character_coverage={character_coverage}",
              f"--model_prefix={model_fp.name}",
              f"--model_type={model_type}",
          ]
      )
      SentencePieceTrainer.Train(argstr)
      if jax.process_index() == 0:
          # Use an intermediate filename that is renamed to the target name to address
          # create and fill delays.
          copy_rename_path = abs_model_path + ".rntmp"
          tf.io.gfile.copy(
              model_fp.name + ".model", copy_rename_path, overwrite=True
          )
          tf.io.gfile.rename(
              copy_rename_path, abs_model_path, overwrite=True
          )
          logging.info(
              "copied %s to %s",
              model_fp.name + ".model",
              abs_model_path,
          )
      else:
          while not tf.io.gfile.exists(abs_model_path):
              time.sleep(1)
          time.sleep(1)
      return abs_model_path
  
file_path: flax/flax/experimental/nnx/examples/lm1b/tokenizer.py
imports: import dataclasses
  import os
  import tempfile
  import time
  from typing import Any, Dict, Iterable, Tuple
  import jax
  import tensorflow
  import tensorflow_text
  from absl import logging
  from sentencepiece import SentencePieceTrainer
label_desc: Use image node_id 2 for calling a global function with example usage: _train_sentencepiece(dataset) and returns: abs_model_path