# Path: t5/models/mesh_transformer.py
# DEPRECATED_GIN_REFERENCES = (
#     "configurable_vocabulary",
#     "get_sentencepiece_model_path",
#     "maybe_print_dataset",
#     "num_parallel_calls",
#     "SentencePieceVocabulary",
#     "t5.data.sentencepiece_vocabulary.SentencePieceVocabulary",
#     "t5.models.mesh_transformer.get_sentencepiece_model_path",
#     "train_model",
#     "vocabularies.Vocabulary",
#     "Vocabulary",
# )
# def mesh_train_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     vocabulary=None,
#     dataset_split=tfds.Split.TRAIN,
#     shuffle=True,
#     seed=None,
#     use_cached=False,
#     pack=True):
#   def _filter_features(ex):
# def mesh_inference_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     shuffle=False,
#     seed=None,
#     vocabulary=None,
#     num_inference_examples=-1,
#     use_cached=False,
#     priming_sequence_length=None):
#   def _split_targets_for_primed_inference(ex):
#   def _prepare_for_unprimed_inference(ex):
#   def _get_dataset_for_single_task(task, sequence_length):
# def mesh_eval_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     vocabulary=None,
#     num_eval_examples=-1,
#     use_cached=False,
#     pack=False,
#     shuffle_eval_examples=False,
#     seed=None):
#   def _get_dataset_for_single_task(task, sequence_length):
# def tsv_dataset_fn(
#     filename,
#     sequence_length,
#     dataset_split,
#     vocabulary,
#     shuffle_buffer_size=10000):
# def get_vocabulary(mixture_or_task_name=None):
# 
# Path: t5/models/utils.py
# _MODEL_FEATURES = [
#     "inputs", "inputs_position", "inputs_segmentation", "targets",
#     "targets_position", "targets_segmentation", "targets_subsegmentation"
# ]
# def filter_features(ex):
# def write_lines_to_file(lines, filename):
# def get_step_from_checkpoint_path(checkpoint_path):
# def write_targets_and_examples(summary_dir, targets, datasets):
# def get_vocabulary(mixture_or_task_name=None):
# def get_latest_checkpoint_from_dir(model_dir):
# def get_checkpoints_iterator(checkpoint_steps, model_dir):
#   def _get_closest_checkpoint(target_checkpoint):
#     def _generate_checkpoints():
#   def __call__(
#       self,
#       checkpoint_step: int,
#       vocabulary: seqio.Vocabulary,
#       tasks: Sequence[seqio.Task],
#       datasets: Mapping[str, tf.data.Dataset],
#       sequence_length: Union[None, Mapping[str, int]]
#   ) -> MutableSequence[Union[str, float]]: ...
#   def __call__(
#       self,
#       task: seqio.Task,
#       sequence_length: Mapping[str, int],
#       split: str,
#   ) -> tf.data.Dataset: ...
# def run_eval(
#     mixture_or_task_name: str,
#     predict_or_score_fn: PredictOrScoreFnCallable,
#     checkpoint_steps: Iterable[int],
#     dataset_fn: DatasetFnCallable,
#     summary_dir: Optional[str] = None,
#     split: Optional[str] = "validation",
#     sequence_length: Optional[Mapping[str, int]] = None,
#     batch_size: Optional[int] = None):
# class PredictOrScoreFnCallable(typing_extensions.Protocol):
# class DatasetFnCallable(typing_extensions.Protocol):
# 
# Path: t5/models/t5_model.py
# class T5Model(metaclass=abc.ABCMeta):
#   """Abstract Base class for T5 Model API."""
# 
#   @abc.abstractmethod
#   def train(self, mixture_or_task_name, steps):
#     raise NotImplementedError()
# 
#   @abc.abstractmethod
#   def eval(self, mixture_or_task_name):
#     raise NotImplementedError()
# 
#   @abc.abstractmethod
#   def predict(self):
#     raise NotImplementedError()
# 
#   @abc.abstractmethod
#   def finetune(self, mixture_or_task_name, finetune_steps):
#     raise NotImplementedError()


import functools
import os
import gin
import gin.tf
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow import optimize
from mesh_tensorflow.transformer import dataset as transformer_dataset
from mesh_tensorflow.transformer import learning_rate_schedules
from mesh_tensorflow.transformer import utils as mtf_utils
from t5.models import mesh_transformer
from t5.models import utils
from t5.models.t5_model import T5Model
                           checkpoint_step,
                           sequence_length,
                           split,
                           eval_with_score=False,
                           **unused_kwargs):
    """Helper function used by eval method to generate predictions or scores.

    Args:
      tasks: list, list of valid tasks to generate predictions or scores.
      vocabulary: a t5.data.vocabulary object or a tuple with separate
        vocabularies for inputs and targets,
      checkpoint_step: integer, step to evaluate the tasks.
      sequence_length: a dict, dictionary with sequence length for inputs and
        targets.
      split: string, split to run the evaluation on.
      eval_with_score: bool, whether to compute log likelihood of targets
        instead of predictions.
    Returns:
      list of decoded predictions or scores depending on eval_with_score flag.
    """
    estimator = self.estimator(
        vocabulary, score_in_predict_mode=eval_with_score,
        sequence_length=sequence_length)

    def estimator_input_fn(params):
      """Eval input function for estimator."""
      del params
      # Concatenate all dataset inputs to only have to do one decode loop
      combined_ds = None
      for task in tasks:

# Path: t5/models/mesh_transformer.py
# DEPRECATED_GIN_REFERENCES = (
#     "configurable_vocabulary",
#     "get_sentencepiece_model_path",
#     "maybe_print_dataset",
#     "num_parallel_calls",
#     "SentencePieceVocabulary",
#     "t5.data.sentencepiece_vocabulary.SentencePieceVocabulary",
#     "t5.models.mesh_transformer.get_sentencepiece_model_path",
#     "train_model",
#     "vocabularies.Vocabulary",
#     "Vocabulary",
# )
# def mesh_train_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     vocabulary=None,
#     dataset_split=tfds.Split.TRAIN,
#     shuffle=True,
#     seed=None,
#     use_cached=False,
#     pack=True):
#   def _filter_features(ex):
# def mesh_inference_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     shuffle=False,
#     seed=None,
#     vocabulary=None,
#     num_inference_examples=-1,
#     use_cached=False,
#     priming_sequence_length=None):
#   def _split_targets_for_primed_inference(ex):
#   def _prepare_for_unprimed_inference(ex):
#   def _get_dataset_for_single_task(task, sequence_length):
# def mesh_eval_dataset_fn(
#     mixture_or_task_name,
#     sequence_length,
#     dataset_split,
#     vocabulary=None,
#     num_eval_examples=-1,
#     use_cached=False,
#     pack=False,
#     shuffle_eval_examples=False,
#     seed=None):
#   def _get_dataset_for_single_task(task, sequence_length):
# def tsv_dataset_fn(
#     filename,
#     sequence_length,
#     dataset_split,
#     vocabulary,
#     shuffle_buffer_size=10000):
# def get_vocabulary(mixture_or_task_name=None):
# 
# Path: t5/models/utils.py
# _MODEL_FEATURES = [
#     "inputs", "inputs_position", "inputs_segmentation", "targets",
#     "targets_position", "targets_segmentation", "targets_subsegmentation"
# ]
# def filter_features(ex):
# def write_lines_to_file(lines, filename):
# def get_step_from_checkpoint_path(checkpoint_path):
# def write_targets_and_examples(summary_dir, targets, datasets):
# def get_vocabulary(mixture_or_task_name=None):
# def get_latest_checkpoint_from_dir(model_dir):
# def get_checkpoints_iterator(checkpoint_steps, model_dir):
#   def _get_closest_checkpoint(target_checkpoint):
#     def _generate_checkpoints():
#   def __call__(
#       self,
#       checkpoint_step: int,
#       vocabulary: seqio.Vocabulary,
#       tasks: Sequence[seqio.Task],
#       datasets: Mapping[str, tf.data.Dataset],
#       sequence_length: Union[None, Mapping[str, int]]
#   ) -> MutableSequence[Union[str, float]]: ...
#   def __call__(
#       self,
#       task: seqio.Task,
#       sequence_length: Mapping[str, int],
#       split: str,
#   ) -> tf.data.Dataset: ...
# def run_eval(
#     mixture_or_task_name: str,
#     predict_or_score_fn: PredictOrScoreFnCallable,
#     checkpoint_steps: Iterable[int],
#     dataset_fn: DatasetFnCallable,
#     summary_dir: Optional[str] = None,
#     split: Optional[str] = "validation",
#     sequence_length: Optional[Mapping[str, int]] = None,
#     batch_size: Optional[int] = None):
# class PredictOrScoreFnCallable(typing_extensions.Protocol):
# class DatasetFnCallable(typing_extensions.Protocol):
# 
# Path: t5/models/t5_model.py
# class T5Model(metaclass=abc.ABCMeta):
#   """Abstract Base class for T5 Model API."""
# 
#   @abc.abstractmethod
#   def train(self, mixture_or_task_name, steps):
#     raise NotImplementedError()
# 
#   @abc.abstractmethod
#   def eval(self, mixture_or_task_name):
#     raise NotImplementedError()
# 
#   @abc.abstractmethod
#   def predict(self):
#     raise NotImplementedError()
# 
#   @abc.abstractmethod
#   def finetune(self, mixture_or_task_name, finetune_steps):
#     raise NotImplementedError()


# Path: t5/models/mtf_model.py
import functools
import os
import gin
import gin.tf
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow import optimize
from mesh_tensorflow.transformer import dataset as transformer_dataset
from mesh_tensorflow.transformer import learning_rate_schedules
from mesh_tensorflow.transformer import utils as mtf_utils
from t5.models import mesh_transformer
from t5.models import utils
from t5.models.t5_model import T5Model

                           checkpoint_step,
                           sequence_length,
                           split,
                           eval_with_score=False,
                           **unused_kwargs):
    """Helper function used by eval method to generate predictions or scores.

    Args:
      tasks: list, list of valid tasks to generate predictions or scores.
      vocabulary: a t5.data.vocabulary object or a tuple with separate
        vocabularies for inputs and targets,
      checkpoint_step: integer, step to evaluate the tasks.
      sequence_length: a dict, dictionary with sequence length for inputs and
        targets.
      split: string, split to run the evaluation on.
      eval_with_score: bool, whether to compute log likelihood of targets
        instead of predictions.
    Returns:
      list of decoded predictions or scores depending on eval_with_score flag.
    """
    estimator = self.estimator(
        vocabulary, score_in_predict_mode=eval_with_score,
        sequence_length=sequence_length)

    def estimator_input_fn(params):
      """Eval input function for estimator."""
      del params
      # Concatenate all dataset inputs to only have to do one decode loop
      combined_ds = None
      for task in tasks:

        ds = mesh_transformer.mesh_eval_dataset_fn(