# 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
        functools.partial(self.flax_module.init, train=False))
    xs = [jnp.ones(inputs.shape), jnp.ones(targets_shape)]
    init_dict = model_init_fn({
        'params': params_rng,
        'dropout': dropout_rng
    }, *xs)
    return init_dict['cache']

  def decode_tokens(self, toks):
    valid_toks = toks[:np.argmax(toks == self.eos_id) + 1].astype(np.int32)
    return self.encoder.detokenize(valid_toks).numpy().decode('utf-8')

  def current_batch_size(self, batch):
    # we assume first token is non-zero in each target side example.
    return int(batch['weights'][:, 0].sum())

  def build_predictor(self):
    decoder = functools.partial(
        decode.decode_step,
        eos_id=self.eos_id,
        beam_size=self.beam_size)
    self.pmapped_predictor = jax.pmap(
        decoder, axis_name='batch', static_broadcasted_argnums=(1, 4))

  def translate_and_calculate_bleu(self):
    """Iterate over all checkpoints and calculate BLEU."""
    self.build_predictor()
    bleu_scores_list = []
    for _, step in self.iterate_checkpoints():
      # If already done, don't redo.

# 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

        functools.partial(self.flax_module.init, train=False))
    xs = [jnp.ones(inputs.shape), jnp.ones(targets_shape)]
    init_dict = model_init_fn({
        'params': params_rng,
        'dropout': dropout_rng
    }, *xs)
    return init_dict['cache']

  def decode_tokens(self, toks):
    valid_toks = toks[:np.argmax(toks == self.eos_id) + 1].astype(np.int32)
    return self.encoder.detokenize(valid_toks).numpy().decode('utf-8')

  def current_batch_size(self, batch):
    # we assume first token is non-zero in each target side example.
    return int(batch['weights'][:, 0].sum())

  def build_predictor(self):
    decoder = functools.partial(
        decode.decode_step,
        eos_id=self.eos_id,
        beam_size=self.beam_size)
    self.pmapped_predictor = jax.pmap(
        decoder, axis_name='batch', static_broadcasted_argnums=(1, 4))

  def translate_and_calculate_bleu(self):
    """Iterate over all checkpoints and calculate BLEU."""
    self.build_predictor()
    bleu_scores_list = []
    for _, step in self.iterate_checkpoints():
      # If already done, don't redo.

      bleu_score = eval_utils.load_evals(self.checkpoint_dir, step,