# Path: prompt_tuning/data/features.py
# T5_FEATURES = t5_tasks.DEFAULT_OUTPUT_FEATURES
# MODEL_TO_FEATURES = {
#     # We use the keys to prefix the tasks t5 is the default so we don't want a
#     # prefix, thus the strange empty string as the key.
#     "": T5_FEATURES,
# }
#   MT5_FEATURES = mt5_tasks.DEFAULT_OUTPUT_FEATURES
#   BYT5_FEATURES = byt5_tasks.DEFAULT_BYTE_OUTPUT_FEATURES
# 
# Path: prompt_tuning/data/metrics.py
# def metric_with_examples(metric_fn,
#                          targets,
#                          predictions,
#                          prediction_field: str = constants.PREDICTION):
# def format_classification(indices: Sequence[int],
#                           examples: List[Mapping[str, Any]],
#                           input_field: str = constants.INPUT_TEXT,
#                           target_field: str = constants.TARGET_TEXT,
#                           prediction_field: str = constants.PREDICTION_TEXT,
#                           **kwargs) -> str:
# def format_qa(indices: Sequence[int],
#               examples: List[Mapping[str, Any]],
#               context_field: str = constants.CONTEXT_TEXT,
#               question_field: str = constants.QUESTION_TEXT,
#               answers_field: str = constants.ANSWERS_TEXT,
#               prediction_field: str = constants.PREDICTION_TEXT,
#               **kwargs) -> str:
# def label_set_stats(
#     targets,
#     predictions,
#     task_name: str,
#     target_field: str = constants.TARGET_TEXT,
#     prediction_field: str = constants.PREDICTION_TEXT,
#     display_sets: bool = False,
# ) -> Dict[str, seqio.metrics.Text]:
# def safe_sample(num_examples: int,
#                 population: List[Mapping[str, Any]],
#                 seed: Optional[int] = None) -> Sequence[int]:
# def text_examples(
#     targets,  # pylint: disable=unused-argument
#     predictions,
#     task_name: str,
#     num_examples: int = 8,
#     seed: Optional[int] = None,
#     find_negative: bool = True,
#     format_fn=format_classification,
#     input_field: str = constants.INPUT_TEXT,
#     target_field: str = constants.TARGET_TEXT,
#     prediction_field: str = constants.PREDICTION,
#     prediction_text_field: str = constants.PREDICTION_TEXT,
#     context_field: str = constants.CONTEXT_TEXT,
#     question_field: str = constants.QUESTION_TEXT,
#     answers_field: str = constants.ANSWERS_TEXT,
# ) -> Dict[str, seqio.metrics.MetricValue]:
# 
# Path: prompt_tuning/data/postprocessors.py
# def remove_extra_id_0(decoded_model_output: str, *args, **kwargs) -> str:  # pylint: disable=unused-argument
# def postprocess_with_examples(
#     postprocess_func,
#     decoded_model_output,
#     prediction_field: str = constants.PREDICTION,
#     prediction_text_field: str = constants.PREDICTION_TEXT,
#     example_fields: Sequence[str] = (constants.INPUT_TEXT,
#                                      constants.TARGET_TEXT),
#     decoded_model_output_field: Optional[str] = None,
#     **kwargs,
# ) -> Union[Mapping[str, Any], Any]:
# def sequential(*funcs):
#   def execute(x, *args, **kwargs):
# def mrpc_to_qqp(decoded_model_output: str, *args, **kwargs) -> str:  # pylint: disable=unused-argument
# def qqp_to_mrpc(decoded_model_output: str, *args, **kwargs) -> str:  # pylint: disable=unused-argument
# 
# Path: prompt_tuning/data/utils.py
# def identity(x: Any, *args, **kwargs) -> Any:  # pylint: disable=unused-argument
# def task_mapping(tasks: Sequence[str],
#                  aliases: Optional[Mapping[str, str]]) -> Dict[str, int]:
# def remove_prefix(s: str, prefix: str) -> str:
# def remove_suffix(string: str, suffix: str) -> str:


import functools
import seqio
from prompt_tuning.data import features
from prompt_tuning.data import metrics as rom_metrics
from prompt_tuning.data import postprocessors as rom_postprocessors
from prompt_tuning.data import utils
from t5.data import preprocessors
from t5.evaluation import metrics
              preprocessors.summarize,
              article_key="document",
              summary_key="summary"),
          seqio.preprocessors.tokenize,
          seqio.CacheDatasetPlaceholder(),
          seqio.preprocessors.append_eos_after_trim,
      ],
      postprocess_fn=functools.partial(
          rom_postprocessors.postprocess_with_examples, utils.identity),
      metric_fns=[
          functools.partial(rom_metrics.metric_with_examples, metrics.rouge),
          functools.partial(
              rom_metrics.text_examples, task_name="xsum", find_negative=False),
      ],
      output_features=feats)

  # ===== CNN/DailyMail =====
  seqio.TaskRegistry.add(
      f"{model_prefix}cnn_dailymail_v310_examples",
      source=seqio.TfdsDataSource(tfds_name="cnn_dailymail:3.1.0"),
      preprocessors=[
          functools.partial(
              preprocessors.summarize,
              article_key="article",
              summary_key="highlights"),
          seqio.preprocessors.tokenize,
          seqio.CacheDatasetPlaceholder(),
          seqio.preprocessors.append_eos_after_trim,
      ],
      postprocess_fn=functools.partial(

# Path: prompt_tuning/data/features.py
# T5_FEATURES = t5_tasks.DEFAULT_OUTPUT_FEATURES
# MODEL_TO_FEATURES = {
#     # We use the keys to prefix the tasks t5 is the default so we don't want a
#     # prefix, thus the strange empty string as the key.
#     "": T5_FEATURES,
# }
#   MT5_FEATURES = mt5_tasks.DEFAULT_OUTPUT_FEATURES
#   BYT5_FEATURES = byt5_tasks.DEFAULT_BYTE_OUTPUT_FEATURES
# 
# Path: prompt_tuning/data/metrics.py
# def metric_with_examples(metric_fn,
#                          targets,
#                          predictions,
#                          prediction_field: str = constants.PREDICTION):
# def format_classification(indices: Sequence[int],
#                           examples: List[Mapping[str, Any]],
#                           input_field: str = constants.INPUT_TEXT,
#                           target_field: str = constants.TARGET_TEXT,
#                           prediction_field: str = constants.PREDICTION_TEXT,
#                           **kwargs) -> str:
# def format_qa(indices: Sequence[int],
#               examples: List[Mapping[str, Any]],
#               context_field: str = constants.CONTEXT_TEXT,
#               question_field: str = constants.QUESTION_TEXT,
#               answers_field: str = constants.ANSWERS_TEXT,
#               prediction_field: str = constants.PREDICTION_TEXT,
#               **kwargs) -> str:
# def label_set_stats(
#     targets,
#     predictions,
#     task_name: str,
#     target_field: str = constants.TARGET_TEXT,
#     prediction_field: str = constants.PREDICTION_TEXT,
#     display_sets: bool = False,
# ) -> Dict[str, seqio.metrics.Text]:
# def safe_sample(num_examples: int,
#                 population: List[Mapping[str, Any]],
#                 seed: Optional[int] = None) -> Sequence[int]:
# def text_examples(
#     targets,  # pylint: disable=unused-argument
#     predictions,
#     task_name: str,
#     num_examples: int = 8,
#     seed: Optional[int] = None,
#     find_negative: bool = True,
#     format_fn=format_classification,
#     input_field: str = constants.INPUT_TEXT,
#     target_field: str = constants.TARGET_TEXT,
#     prediction_field: str = constants.PREDICTION,
#     prediction_text_field: str = constants.PREDICTION_TEXT,
#     context_field: str = constants.CONTEXT_TEXT,
#     question_field: str = constants.QUESTION_TEXT,
#     answers_field: str = constants.ANSWERS_TEXT,
# ) -> Dict[str, seqio.metrics.MetricValue]:
# 
# Path: prompt_tuning/data/postprocessors.py
# def remove_extra_id_0(decoded_model_output: str, *args, **kwargs) -> str:  # pylint: disable=unused-argument
# def postprocess_with_examples(
#     postprocess_func,
#     decoded_model_output,
#     prediction_field: str = constants.PREDICTION,
#     prediction_text_field: str = constants.PREDICTION_TEXT,
#     example_fields: Sequence[str] = (constants.INPUT_TEXT,
#                                      constants.TARGET_TEXT),
#     decoded_model_output_field: Optional[str] = None,
#     **kwargs,
# ) -> Union[Mapping[str, Any], Any]:
# def sequential(*funcs):
#   def execute(x, *args, **kwargs):
# def mrpc_to_qqp(decoded_model_output: str, *args, **kwargs) -> str:  # pylint: disable=unused-argument
# def qqp_to_mrpc(decoded_model_output: str, *args, **kwargs) -> str:  # pylint: disable=unused-argument
# 
# Path: prompt_tuning/data/utils.py
# def identity(x: Any, *args, **kwargs) -> Any:  # pylint: disable=unused-argument
# def task_mapping(tasks: Sequence[str],
#                  aliases: Optional[Mapping[str, str]]) -> Dict[str, int]:
# def remove_prefix(s: str, prefix: str) -> str:
# def remove_suffix(string: str, suffix: str) -> str:


# Path: prompt_tuning/data/summarization.py
import functools
import seqio
from prompt_tuning.data import features
from prompt_tuning.data import metrics as rom_metrics
from prompt_tuning.data import postprocessors as rom_postprocessors
from prompt_tuning.data import utils
from t5.data import preprocessors
from t5.evaluation import metrics

              preprocessors.summarize,
              article_key="document",
              summary_key="summary"),
          seqio.preprocessors.tokenize,
          seqio.CacheDatasetPlaceholder(),
          seqio.preprocessors.append_eos_after_trim,
      ],
      postprocess_fn=functools.partial(
          rom_postprocessors.postprocess_with_examples, utils.identity),
      metric_fns=[
          functools.partial(rom_metrics.metric_with_examples, metrics.rouge),
          functools.partial(
              rom_metrics.text_examples, task_name="xsum", find_negative=False),
      ],
      output_features=feats)

  # ===== CNN/DailyMail =====
  seqio.TaskRegistry.add(
      f"{model_prefix}cnn_dailymail_v310_examples",
      source=seqio.TfdsDataSource(tfds_name="cnn_dailymail:3.1.0"),
      preprocessors=[
          functools.partial(
              preprocessors.summarize,
              article_key="article",
              summary_key="highlights"),
          seqio.preprocessors.tokenize,
          seqio.CacheDatasetPlaceholder(),
          seqio.preprocessors.append_eos_after_trim,
      ],
      postprocess_fn=functools.partial(

          rom_postprocessors.postprocess_with_examples, utils.identity),