code: def load_or_train_tokenizer(
      dataset: tf.data.Dataset,
      *,
      vocab_path: str,
      vocab_size: int,
      max_corpus_chars: int,
      data_keys: Tuple[str, str] = ("inputs", "targets"),
  ):
      """Loads the tokenizer at `vocab_path` or trains a one from `dataset`."""
      try:
          return _load_sentencepiece_tokenizer(vocab_path)
      except tf.errors.NotFoundError:
          logging.info(
              "SentencePiece vocab not found, building one from data."
          )
          vocab_path = _train_sentencepiece(
              dataset,
              vocab_size=vocab_size,
              maxchars=max_corpus_chars,
              model_path=vocab_path,
              data_keys=data_keys,
          )
          return _load_sentencepiece_tokenizer(vocab_path)
  
file_path: flax/examples/wmt/tokenizer.py
imports: import dataclasses
  import os
  import tempfile
  import time
  from typing import Any, Dict, Iterable, Tuple
  from absl import logging
  import jax
  import tensorflow
  import tensorflow_text
  from sentencepiece import SentencePieceTrainer
label_desc: Use image node_id 4 for calling a global function with example usage: load_or_train_tokenizer(dataset) and returns: _load_sentencepiece_tokenizer, _load_sentencepiece_tokenizer