# 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
        {'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)
    optimizer_state = optimizer_init_fn(params)
    self.optimizer_state = jax_utils.replicate(optimizer_state)
    self.pmapped_init_cache = jax.pmap(
        functools.partial(
            self.initialize_cache,
            max_length=self.max_length,
            params_rng=params_rng,
            dropout_rng=dropout_rng))

  def initialize_cache(self, inputs, max_length, params_rng, dropout_rng):
    """Initialize a cache for a given input shape and max decode length."""
    targets_shape = (inputs.shape[0], max_length) + inputs.shape[2:]
    model_init_fn = jax.jit(
        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)

# 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

        {'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)
    optimizer_state = optimizer_init_fn(params)
    self.optimizer_state = jax_utils.replicate(optimizer_state)
    self.pmapped_init_cache = jax.pmap(
        functools.partial(
            self.initialize_cache,
            max_length=self.max_length,
            params_rng=params_rng,
            dropout_rng=dropout_rng))

  def initialize_cache(self, inputs, max_length, params_rng, dropout_rng):
    """Initialize a cache for a given input shape and max decode length."""
    targets_shape = (inputs.shape[0], max_length) + inputs.shape[2:]
    model_init_fn = jax.jit(
        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')