# Path: init2winit/dataset_lib/lm1b_tokenizer.py
# def _dump_chars_to_textfile(
#     dataset: tf.data.Dataset,
#     maxchars: int = int(1e7),
#     data_keys=('inputs', 'targets')
# ) -> Tuple[str, int]:
# 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')):
# def _load_sentencepiece_tokenizer(model_path: str,
#                                   add_bos: bool = False,
#                                   add_eos: bool = True,
#                                   reverse: bool = False):
# 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')):
#   def __call__(self, features: Features) -> Features:
# class TokenizeOp:


import os
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import Dict, List, Optional, Union
from absl import logging
from clu import deterministic_data
from init2winit.dataset_lib import lm1b_tokenizer

  if prefetch_size:
    dataset = dataset.prefetch(prefetch_size)

  return dataset


def get_lm1b_datasets(hps, per_host_batch_size, per_host_eval_batch_size,
                      shuffle_rng):
  """Load and return dataset of batched examples for use during training."""
  if hps.vocab_path is None:
    vocab_path = os.path.expanduser('~/lm2b_sentencepiece_model')
  else:
    vocab_path = hps.vocab_path

  logging.info('vocab_path = %s', vocab_path)

  train_ds_builder = tfds.builder('lm1b')
  train_data = get_raw_dataset(train_ds_builder, 'train')

  eval_ds_builder = train_ds_builder
  eval_data = get_raw_dataset(eval_ds_builder, hps.eval_split)

  # Tokenize data.
  sp_tokenizer = lm1b_tokenizer.load_or_train_tokenizer(
      train_data,
      vocab_path=vocab_path,
      vocab_size=hps.vocab_size,
      max_corpus_chars=hps.max_corpus_chars)
  train_data = train_data.map(

# Path: init2winit/dataset_lib/lm1b_tokenizer.py
# def _dump_chars_to_textfile(
#     dataset: tf.data.Dataset,
#     maxchars: int = int(1e7),
#     data_keys=('inputs', 'targets')
# ) -> Tuple[str, int]:
# 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')):
# def _load_sentencepiece_tokenizer(model_path: str,
#                                   add_bos: bool = False,
#                                   add_eos: bool = True,
#                                   reverse: bool = False):
# 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')):
#   def __call__(self, features: Features) -> Features:
# class TokenizeOp:


# Path: init2winit/dataset_lib/lm1b_input_pipeline_v2.py
import os
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import Dict, List, Optional, Union
from absl import logging
from clu import deterministic_data
from init2winit.dataset_lib import lm1b_tokenizer


  if prefetch_size:
    dataset = dataset.prefetch(prefetch_size)

  return dataset


def get_lm1b_datasets(hps, per_host_batch_size, per_host_eval_batch_size,
                      shuffle_rng):
  """Load and return dataset of batched examples for use during training."""
  if hps.vocab_path is None:
    vocab_path = os.path.expanduser('~/lm2b_sentencepiece_model')
  else:
    vocab_path = hps.vocab_path

  logging.info('vocab_path = %s', vocab_path)

  train_ds_builder = tfds.builder('lm1b')
  train_data = get_raw_dataset(train_ds_builder, 'train')

  eval_ds_builder = train_ds_builder
  eval_data = get_raw_dataset(eval_ds_builder, hps.eval_split)

  # Tokenize data.
  sp_tokenizer = lm1b_tokenizer.load_or_train_tokenizer(
      train_data,
      vocab_path=vocab_path,
      vocab_size=hps.vocab_size,
      max_corpus_chars=hps.max_corpus_chars)
  train_data = train_data.map(

      lm1b_tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE)