# 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/mt_eval/decode.py
# EOS_ID = 2
# NEG_INF = np.array(-1.0e7)
# def brevity_penalty(alpha, length):
# def add_beam_dim(x, beam_size):
# def flatten_beam_dim(x):
# def unflatten_beam_dim(x, batch_size, beam_size):
# def flat_batch_beam_expand(x, beam_size):
# def gather_beams(nested, beam_indices, batch_size, new_beam_size):
#   def gather_fn(x):
# def gather_topk_beams(nested, score_or_log_prob, batch_size, new_beam_size):
# def beam_init(batch_size, beam_size, max_decode_len, cache):
# def beam_search(inputs,
#                 cache,
#                 tokens_to_logits,
#                 beam_size=4,
#                 alpha=0.6,
#                 eos_id=EOS_ID,
#                 max_decode_len=None):
#   def beam_search_loop_cond_fn(state):
#   def beam_search_loop_body_fn(state):
# def decode_step(batch,
#                 flax_module,
#                 params,
#                 cache,
#                 max_decode_len,
#                 eos_id=EOS_ID,
#                 beam_size=4):
#   def tokens_ids_to_logits(flat_ids, flat_cache):
# class BeamState:
# 
# Path: init2winit/mt_eval/eval_utils.py
# def compute_bleu_from_predictions(predictions, references, name):
# def get_eval_fpath(ckpt_dir, ckpt_step, eval_split):
# def load_evals(ckpt_dir, ckpt_step, eval_split):
# def save_evals(ckpt_dir, ckpt_step, eval_split, bleu_score):
# def _load_checkpoint(checkpoint_path, params, optimizer_state, batch_stats,
#                      replicate=True):
# def average_checkpoints(
#     checkpoint_paths, params, optimizer_state, batch_stats):
# def get_checkpoints_in_range(checkpoint_dir, lower_bound, upper_bound):
# 
# Path: init2winit/optimizer_lib/optimizers.py
# def sgd(learning_rate, weight_decay, momentum=None, nesterov=False):
# def get_optimizer(hps, model=None):
# def _wrap_update_fn(opt_name, opt_update):
#   def update_fn(grads, optimizer_state, params, batch=None, batch_stats=None):


import copy
import functools
import os
import jax
import jax.numpy as jnp
import numpy as np
from absl import logging
from flax import jax_utils
from flax.training import common_utils
from init2winit.dataset_lib import mt_tokenizer
from init2winit.mt_eval import decode
from init2winit.mt_eval import eval_utils
from init2winit.optimizer_lib import optimizers
from tensorflow.io import gfile
    else:
      logging.info('Loading test split')
      return self.dataset.test_epoch(self.eval_num_batches)

  def load_tokenizer(self, vocab_path):
    logging.info('Loading tokenizer from: %s', vocab_path)
    # pylint: disable=protected-access
    return mt_tokenizer._load_sentencepiece_tokenizer(vocab_path)
    # pylint: enable=protected-access

  def initialize_model(self, model_cls, dataset_meta_data, dropout_rng,
                       params_rng):
    """Initialie model, especially cache for fast auto-regressive decoding."""
    loss_name = 'cross_entropy'
    metrics_name = 'classification_metrics'
    hps = copy.deepcopy(self.hps)
    hps = hps.unlock()
    hps.decode = True
    hps = hps.lock()
    model = model_cls(hps, dataset_meta_data, loss_name, metrics_name)
    xs = [np.zeros((self.hps.batch_size, *x)) for x in self.hps.input_shape]
    model_init_fn = jax.jit(
        functools.partial(model.flax_module.init, train=False))
    init_dict = model_init_fn(
        {'params': params_rng, 'dropout': dropout_rng}, *xs)
    params = init_dict['params']
    batch_stats = init_dict.get('batch_stats', {})
    self.flax_module = model.flax_module
    self.params = params
    self.batch_stats = jax_utils.replicate(batch_stats)

# 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/mt_eval/decode.py
# EOS_ID = 2
# NEG_INF = np.array(-1.0e7)
# def brevity_penalty(alpha, length):
# def add_beam_dim(x, beam_size):
# def flatten_beam_dim(x):
# def unflatten_beam_dim(x, batch_size, beam_size):
# def flat_batch_beam_expand(x, beam_size):
# def gather_beams(nested, beam_indices, batch_size, new_beam_size):
#   def gather_fn(x):
# def gather_topk_beams(nested, score_or_log_prob, batch_size, new_beam_size):
# def beam_init(batch_size, beam_size, max_decode_len, cache):
# def beam_search(inputs,
#                 cache,
#                 tokens_to_logits,
#                 beam_size=4,
#                 alpha=0.6,
#                 eos_id=EOS_ID,
#                 max_decode_len=None):
#   def beam_search_loop_cond_fn(state):
#   def beam_search_loop_body_fn(state):
# def decode_step(batch,
#                 flax_module,
#                 params,
#                 cache,
#                 max_decode_len,
#                 eos_id=EOS_ID,
#                 beam_size=4):
#   def tokens_ids_to_logits(flat_ids, flat_cache):
# class BeamState:
# 
# Path: init2winit/mt_eval/eval_utils.py
# def compute_bleu_from_predictions(predictions, references, name):
# def get_eval_fpath(ckpt_dir, ckpt_step, eval_split):
# def load_evals(ckpt_dir, ckpt_step, eval_split):
# def save_evals(ckpt_dir, ckpt_step, eval_split, bleu_score):
# def _load_checkpoint(checkpoint_path, params, optimizer_state, batch_stats,
#                      replicate=True):
# def average_checkpoints(
#     checkpoint_paths, params, optimizer_state, batch_stats):
# def get_checkpoints_in_range(checkpoint_dir, lower_bound, upper_bound):
# 
# Path: init2winit/optimizer_lib/optimizers.py
# def sgd(learning_rate, weight_decay, momentum=None, nesterov=False):
# def get_optimizer(hps, model=None):
# def _wrap_update_fn(opt_name, opt_update):
#   def update_fn(grads, optimizer_state, params, batch=None, batch_stats=None):


# Path: init2winit/mt_eval/bleu_evaluator.py
import copy
import functools
import os
import jax
import jax.numpy as jnp
import numpy as np
from absl import logging
from flax import jax_utils
from flax.training import common_utils
from init2winit.dataset_lib import mt_tokenizer
from init2winit.mt_eval import decode
from init2winit.mt_eval import eval_utils
from init2winit.optimizer_lib import optimizers
from tensorflow.io import gfile

    else:
      logging.info('Loading test split')
      return self.dataset.test_epoch(self.eval_num_batches)

  def load_tokenizer(self, vocab_path):
    logging.info('Loading tokenizer from: %s', vocab_path)
    # pylint: disable=protected-access
    return mt_tokenizer._load_sentencepiece_tokenizer(vocab_path)
    # pylint: enable=protected-access

  def initialize_model(self, model_cls, dataset_meta_data, dropout_rng,
                       params_rng):
    """Initialie model, especially cache for fast auto-regressive decoding."""
    loss_name = 'cross_entropy'
    metrics_name = 'classification_metrics'
    hps = copy.deepcopy(self.hps)
    hps = hps.unlock()
    hps.decode = True
    hps = hps.lock()
    model = model_cls(hps, dataset_meta_data, loss_name, metrics_name)
    xs = [np.zeros((self.hps.batch_size, *x)) for x in self.hps.input_shape]
    model_init_fn = jax.jit(
        functools.partial(model.flax_module.init, train=False))
    init_dict = model_init_fn(
        {'params': params_rng, 'dropout': dropout_rng}, *xs)
    params = init_dict['params']
    batch_stats = init_dict.get('batch_stats', {})
    self.flax_module = model.flax_module
    self.params = params
    self.batch_stats = jax_utils.replicate(batch_stats)

    optimizer_init_fn, _ = optimizers.get_optimizer(self.hps)