# Path: prompt_tuning/data/preprocessors.py
# AUTOTUNE = tf.data.experimental.AUTOTUNE
# def preprocess_tsv_to_qa(line,
#                          field_delim='\t',
#                          answer_delim='|||',
#                          field_names: Optional[Sequence[str]] = None):
# def mrqa(example):
# def preprocess_text_generation(
#     example,
#     source_key,
#     target_key,
#     task_name=None,
#     prefix=None,
#     source_nested_key=None,
#     target_nested_key=None,
# ):
# def remove_first_text_token(example, key: str = 'inputs'):
# def add_sentinel_to_targets(ds, output_features):
# def add_sentinel_to_beginning(ds,
#                               output_features,
#                               field: str = 'targets',
#                               offset: int = 0):
#   def add(ex):
# def build_langid_filter(lang_code, lang_detector, threshold):
#   def filter_fn(text):
# def filter_langid(dataset,
#                   lang_code,
#                   lang_detector,
#                   text_key='text',
#                   threshold=0.95):
# def token_deletion(
#     ds: tf.data.Dataset,
#     sequence_length: SequenceLengthType,
#     output_features: seqio.preprocessors.OutputFeaturesType,
#     noise_density: float = 0.15
# ) -> tf.data.Dataset:
# def text_infilling(
#     ds: tf.data.Dataset,
#     sequence_length: SequenceLengthType,
#     output_features: seqio.preprocessors.OutputFeaturesType,
#     noise_density: float = 0.15,
#     mean_noise_span_length: int = 3
# ) -> tf.data.Dataset:


import os
import textwrap
import unittest.mock as mock
import numpy as np
import seqio
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from absl.testing import parameterized
from prompt_tuning.data import preprocessors
from seqio import test_utils

TEST_T5_FEATURES = {
    "inputs": seqio.Feature(
        vocabulary=seqio.SentencePieceVocabulary(
            os.path.join(TEST_DATA, "t5_vocab"), 100),
        add_eos=True,
        required=False),
    "targets": seqio.Feature(
        vocabulary=seqio.SentencePieceVocabulary(
            os.path.join(TEST_DATA, "t5_vocab"), 100),
        add_eos=True)
}


def create_fake_text_dataset(examples: int = 10, text_size: int = TEXT_SIZE):
  text = np.reshape(
      # Start at 2 so we skip EOS=1 which could be a problem on any tests that
      # actually decode the fake inputs.
      np.arange(2, examples * text_size + 2),
      (-1, text_size)).astype(np.int32)
  return tf.data.Dataset.from_tensor_slices({"targets": text})


class PreprocessorsTest(tf.test.TestCase):

  def test_remove_first_text_token(self):
    input_strings = ["This is my first example", "The second"]
    gold_strings = [" ".join(s.split()[1:]) for s in input_strings]
    ds = tf.data.Dataset.from_tensor_slices({"inputs": input_strings})


# Path: prompt_tuning/data/preprocessors.py
# AUTOTUNE = tf.data.experimental.AUTOTUNE
# def preprocess_tsv_to_qa(line,
#                          field_delim='\t',
#                          answer_delim='|||',
#                          field_names: Optional[Sequence[str]] = None):
# def mrqa(example):
# def preprocess_text_generation(
#     example,
#     source_key,
#     target_key,
#     task_name=None,
#     prefix=None,
#     source_nested_key=None,
#     target_nested_key=None,
# ):
# def remove_first_text_token(example, key: str = 'inputs'):
# def add_sentinel_to_targets(ds, output_features):
# def add_sentinel_to_beginning(ds,
#                               output_features,
#                               field: str = 'targets',
#                               offset: int = 0):
#   def add(ex):
# def build_langid_filter(lang_code, lang_detector, threshold):
#   def filter_fn(text):
# def filter_langid(dataset,
#                   lang_code,
#                   lang_detector,
#                   text_key='text',
#                   threshold=0.95):
# def token_deletion(
#     ds: tf.data.Dataset,
#     sequence_length: SequenceLengthType,
#     output_features: seqio.preprocessors.OutputFeaturesType,
#     noise_density: float = 0.15
# ) -> tf.data.Dataset:
# def text_infilling(
#     ds: tf.data.Dataset,
#     sequence_length: SequenceLengthType,
#     output_features: seqio.preprocessors.OutputFeaturesType,
#     noise_density: float = 0.15,
#     mean_noise_span_length: int = 3
# ) -> tf.data.Dataset:


# Path: prompt_tuning/data/preprocessors_test.py
import os
import textwrap
import unittest.mock as mock
import numpy as np
import seqio
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from absl.testing import parameterized
from prompt_tuning.data import preprocessors
from seqio import test_utils


TEST_T5_FEATURES = {
    "inputs": seqio.Feature(
        vocabulary=seqio.SentencePieceVocabulary(
            os.path.join(TEST_DATA, "t5_vocab"), 100),
        add_eos=True,
        required=False),
    "targets": seqio.Feature(
        vocabulary=seqio.SentencePieceVocabulary(
            os.path.join(TEST_DATA, "t5_vocab"), 100),
        add_eos=True)
}


def create_fake_text_dataset(examples: int = 10, text_size: int = TEXT_SIZE):
  text = np.reshape(
      # Start at 2 so we skip EOS=1 which could be a problem on any tests that
      # actually decode the fake inputs.
      np.arange(2, examples * text_size + 2),
      (-1, text_size)).astype(np.int32)
  return tf.data.Dataset.from_tensor_slices({"targets": text})


class PreprocessorsTest(tf.test.TestCase):

  def test_remove_first_text_token(self):
    input_strings = ["This is my first example", "The second"]
    gold_strings = [" ".join(s.split()[1:]) for s in input_strings]
    ds = tf.data.Dataset.from_tensor_slices({"inputs": input_strings})


    processed_ds = preprocessors.remove_first_text_token(ds)