# 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
  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(
      sampled_train_data,
      vocab_path=vocab_path,
      vocab_size=config.vocab_size,
      max_corpus_chars=config.max_corpus_chars,
      character_coverage=config.character_coverage,
      byte_fallback=config.byte_fallback,
      split_digits=config.split_digits,
      user_defined_symbols=user_defined_symbols)
  sampled_train_data = sampled_train_data.map(
      tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE)
  eval_data = eval_data.map(

# 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

  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(
      sampled_train_data,
      vocab_path=vocab_path,
      vocab_size=config.vocab_size,
      max_corpus_chars=config.max_corpus_chars,
      character_coverage=config.character_coverage,
      byte_fallback=config.byte_fallback,
      split_digits=config.split_digits,
      user_defined_symbols=user_defined_symbols)
  sampled_train_data = sampled_train_data.map(
      tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE)
  eval_data = eval_data.map(

      tokenizer.TokenizeOp(sp_tokenizer), num_parallel_calls=AUTOTUNE)