# Path: init2winit/dataset_lib/mt_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,
#                          byte_fallback: bool = False,
#                          split_digits: bool = False,
#                          data_keys: Tuple[str, str] = ('inputs', 'targets'),
#                          user_defined_symbols: List[str] = []):
# 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,
#                             character_coverage: float = 1.0,
#                             byte_fallback: bool = False,
#                             split_digits: bool = False,
#                             data_keys: Tuple[str, str] = ('inputs', 'targets'),
#                             user_defined_symbols: List[str] = []):
#   def __call__(self, features: Features) -> Features:
# class TokenizeOp:


import os
import jax
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import Dict, List, Optional, Union
from clu import deterministic_data
from init2winit.dataset_lib import mt_tokenizer as tokenizer
from ml_collections.config_dict import config_dict
      train_ds_builders, config.train_split, dataset_rates,
      config.reverse_translation, config.add_language_token,
      loss_weights=config.loss_weights,
      is_training=True, sample_seed=sample_seed,
      shuffle_seed=shuffle_seed)

  eval_data = get_raw_dataset(
      eval_ds_builder,
      config.eval_split,
      reverse_translation=config.reverse_translation,
      add_language_token=config.add_language_token)

  predict_data = get_raw_dataset(
      predict_ds_builder,
      config.predict_split,
      reverse_translation=config.reverse_translation,
      add_language_token=config.add_language_token)

  # Tokenize data.
  user_defined_symbols = []
  if config.add_language_token:
    user_defined_symbols = [
        get_user_defined_symbols(ds_builders.info, True)
        for ds_builders in train_ds_builders]
    user_defined_symbols.extend(
        [get_user_defined_symbols(ds_builders.info, False)
         for ds_builders in train_ds_builders])
  user_defined_symbols = list(set(user_defined_symbols))

  # TODO(dxin): Check in vocab file eventually.

# Path: init2winit/dataset_lib/mt_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,
#                          byte_fallback: bool = False,
#                          split_digits: bool = False,
#                          data_keys: Tuple[str, str] = ('inputs', 'targets'),
#                          user_defined_symbols: List[str] = []):
# 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,
#                             character_coverage: float = 1.0,
#                             byte_fallback: bool = False,
#                             split_digits: bool = False,
#                             data_keys: Tuple[str, str] = ('inputs', 'targets'),
#                             user_defined_symbols: List[str] = []):
#   def __call__(self, features: Features) -> Features:
# class TokenizeOp:


# Path: init2winit/dataset_lib/mt_pipeline.py
import os
import jax
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import Dict, List, Optional, Union
from clu import deterministic_data
from init2winit.dataset_lib import mt_tokenizer as tokenizer
from ml_collections.config_dict import config_dict

      train_ds_builders, config.train_split, dataset_rates,
      config.reverse_translation, config.add_language_token,
      loss_weights=config.loss_weights,
      is_training=True, sample_seed=sample_seed,
      shuffle_seed=shuffle_seed)

  eval_data = get_raw_dataset(
      eval_ds_builder,
      config.eval_split,
      reverse_translation=config.reverse_translation,
      add_language_token=config.add_language_token)

  predict_data = get_raw_dataset(
      predict_ds_builder,
      config.predict_split,
      reverse_translation=config.reverse_translation,
      add_language_token=config.add_language_token)

  # Tokenize data.
  user_defined_symbols = []
  if config.add_language_token:
    user_defined_symbols = [
        get_user_defined_symbols(ds_builders.info, True)
        for ds_builders in train_ds_builders]
    user_defined_symbols.extend(
        [get_user_defined_symbols(ds_builders.info, False)
         for ds_builders in train_ds_builders])
  user_defined_symbols = list(set(user_defined_symbols))

  # TODO(dxin): Check in vocab file eventually.

  sp_tokenizer = tokenizer.load_or_train_tokenizer(