Function train_tokenizer

src/MaxText/train_tokenizer.py:123–140

tokenizer training function

(
    dataset: tf.data.Dataset,
    *,
    vocab_path: str,
    vocab_size: int,
    max_corpus_chars: int,
    data_keys: tuple[str] = ("text",),
)

def train_tokenizer(
    dataset: tf.data.Dataset,
    *,
    vocab_path: str,
    vocab_size: int,
    max_corpus_chars: int,
    data_keys: tuple[str] = ("text",),
):
  """tokenizer training function"""
  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,
  )
  logging.info("Model saved at %s", vocab_path)


def main(argv):