<commit_before># SeqIO: Task-based datasets, preprocessing, and evaluation for sequence models.

**SeqIO** is a library for processing sequential data to be fed into downstream
sequence models. It uses [`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset)
to create scalable data pipelines but requires minimal use of TensorFlow. In
particular, with one line of code, the returned dataset can be transformed to a
numpy iterator and hence it is fully compatible with other frameworks such as
[JAX](https://github.com/google/jax) or
[PyTorch](https://pytorch.org/).

Currently, SeqIO assumes that the dataset is a sequence, i.e., each feature is
one-dimensional array. Modalities such as text or audio are naturally supported.
Images are supported as long as they are represented as sequences (e.g.,
[Image GPT](http://proceedings.mlr.press/v119/chen20s.html)). We will release this constraint
in the future in order to support higher dimensional data.

SeqIO is a refactor of the [`t5.data`](https://github.com/google-research/text-to-text-transfer-transformer/)
library used (in conjunction with the [Mesh Tensorflow](https://github.com/tensorflow/mesh)
Transformer implementation) to train the T5 models introduced in
[_Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer_](https://arxiv.org/abs/1910.10683).

If you have used `t5.data` in the past and want to know how SeqIO differs, please read [this section](#differences-from-t5data).

<!--Uncomment once forked.
## Installation

### From Pypi

```sh
pip install seqio
```

### From Source

```sh
git clone https://github.com/google/seqio.git
cd seqio
pip install -e .
-->

## Usage Tutorial

At a high level, we use SeqIO with the following steps.

  1. Define a `Task` (and optionally a `Mixture`).

  1. Define (or use an existing) a `FeatureConverter` based on the model architecture.

  1. Use the top-level function `seqio.get_dataset` to obtain the
     `tf.data.Dataset` instance.

We will look at each of these steps in detail.


### Defining a `Task`

The most important class in SeqIO is the `Task`. It is an abstraction that combines:

  * a raw *data source*
  * one or more *preprocessing* steps
  * a *vocabulary* to tokenize/detokenize each preprocessed feature for the model
  * a *postprocessor* to convert detokenized model outputs into a format for evaluation
  * one or more *metrics* to evaluate with

Oftentimes a `Task` lines up with a common benchmark. In this tutorial, we use
[WMT 19 English-German](http://www.statmt.org/wmt19/translation-task.html) machine
translation task. In the end, our `Task` will look like this:


```py
seqio.TaskRegistry.add(
    "wmt19_ende",
    seqio.TfdsDataSource(tfds_name="wmt19_translate/de-en:1.0.0"),
    preprocessors=[
        functools.partial(
            translate, source_language='en', target_language='de'),
        seqio.preprocessors.tokenize, seqio.preprocessors.append_eos
    ],
    output_features={
        'inputs':
            seqio.Feature(
                seqio.SentencePieceVocabulary('/path/to/inputs/vocab'),
                add_eos=False,
                dtype=tf.int32),
        'targets':
            seqio.Feature(
                seqio.SentencePieceVocabulary('/path/to/targets/vocab'),
                add_eos=True,
                dtype=tf.int32),
    },
    metric_fns=[bleu])
```

We typically add the `Task` to the global registry when we define it (as shown
above) to make it easier to use with model configs and flags. Thus, it  must
have a unique string name (`"wmt19_ende"` in this case). Note, however, that
you may also instantiate a `seqio.Task` directly without adding it to the
registry, if desired.

We'll now break down each part of the task definition.

#### Data Source

Data sources are the first step in your pipeline, providing a way to load raw
data in many formats as a `tf.data.Dataset`.
All data sources are subclasses of the `DataSource` base class and are defined in
[dataset_providers](https://github.com/google/seqio/tree/main/seqio/dataset_providers.py),

Existing implementations include:

  * `TfdsDataSource` for loading examples from [TensorFlow Datasets](https://www.tensorflow.org/datasets).
  * `TextLineDataSource` for loading examples from text files (e.g., tsv).
  * `TFExampleDataSource` for loading [`tf.train.Example`](https://www.tensorflow.org/tutorials/load_data/tfrecord) protos from a file (e.g. a `TFRecord` file.)
  * `FunctionDataSource` for providing an custom function that returns a `tf.data.Dataset`.

In our example, we are using the `TfdsDataSource`. We specify the name of the WMT dataset in TFDS ([`"wmt19_translate"`](https://www.tensorflow.org/datasets/catalog/wmt19_translate)), the specific config for the language pair that excludes the context for the open domain setting (`"de-en"`), and the version number (`"1.0.0"`).

#### Output Features

The `output_features` field expects a dictionary that maps string feature names
to `seqio.Feature` objects. This defines what the `Task` is expected to produce
in its output examples. The output examples *may* contain additional fields, but
they *must* contain these fields in the specified format or exceptions will be
raised.


Each `Feature` includes:

  * A `vocabulary`, which must subclass [`seqio.Vocabulary`](https://github.com/google/seqio/tree/main/seqio/vocabularies.py), to specify how the feature can be tokenized and detokenized. You may use `seqio.PassThroughVocabulary` if tokenization is not necessary.
  * `add_eos`, which specifies whether the feature should end with the vocabulary's EOS token.
  * The output `dtype` which must be a `tf.dtypes.DType`.

**Note:** specifying these options on `Feature` does not by itself ensure the proper transformations are applied -- you must also include the necessary preprocessors.

The [tasks used in T5](TODO) all produce "inputs" and "targets" features to be consumed by the text-to-text model. For a decoder-only language model, only a single feature (e.g., "targets") would be necessary.
Nevertheless, SeqIO is flexible enough to generate arbitrary output features what will be converted into model features by the [`FeatureConverter`](#featureconverter) later in the pipeline.

#### Preprocessors

Preprocessors are functions that transform one `tf.data.Dataset` into a new `tf.data.Dataset`. Typically this involves executing a `map` over the given dataset. The preprocessors provided to the `Task` will be executed sequentially.


As an example, let's look at the previously undefined `translate` from the "wmt19_ende" example above.

```py
def translate(dataset: tf.data.Dataset,
              source_language: str,
              target_language: str) -> tf.data.Dataset:
  def _translate(ex: Mapping[str, tf.Tensor]) -> Mapping[str, tf.Tensor]:
    """Convert a translation example to a text2text pair.

    For example, say the dataset returns examples of this format:
      {'de': 'Das ist gut.', 'en': 'That is good.'}
    If source_language = 'de', target_language = 'en', then the outputs will have
    the format:
      {'inputs': 'translate de to en: Das ist gut.',
      'targets': 'That is good.'}

    Args:
      x: an example to process.
      source_language: source language code (e.g. 'en') to translate from.
      target_language: target language code (e.g. 'de') to translate to.

    Returns:
      A preprocessed example with the format listed above.
    """
    src_str = f'translate {source_language}'
    tgt_str = f' to {target_language}: '
    return {
        'inputs': tf.strings.join([src_str, tgt_str, ex[source_language]]),
        'targets': ex[target_language],
    }

  return dataset.map(_translate,
                     num_parallel_calls=tf.data.experimental.AUTOTUNE)
```

The TFDS dataset provides the dataset where each example has the form: `{'de':
'Das ist gut.', 'en': 'That is good.'}`. We convert this to "inputs" and
"targets" with the appropriate prompt to inform the model of the task.


A few **important** notes:

  1. When instantiating a `Task`, the preprocessor functions can have the following arguments: `dataset`, `output_features`, and `sequence_length`. The first (positional) dataset argument is always required. If an argument named `output_features` is provided, the [output feature mapping](#output-features) will be passed to the preprocessor. If `sequence_length` is provided, a mapping from feature name to its *maximum* final sequence length ([provided by the caller](#getting-a-preprocessed-dataset) will be passed -- any sequences that are too long after preprocessing will be automatically truncated. If a preprocessor function does have other arguments, they must have default values or be bound (e.g., with `functools.partial` as used in `translate`) before instantiating the `Task`.

  1. Mapping functions operate on and return `tf.Tensor`s using TensorFlow operations, although it is possible to take advantage of automatic [AutoGraph](https://blog.tensorflow.org/2018/07/autograph-converts-python-into-tensorflow-graphs.html) conversion for `numpy` or use [`tf.py_function`](https://www.tensorflow.org/api_docs/python/tf/py_function) to wrap arbitrary Python code. See `tf.data.Dataset` [documentation](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) for more details.

  1. When calling `map`, it is important to **always** set `num_parallel_calls=tf.data.experimental.AUTOTUNE` to avoid creating a bottleneck. The `seqio.map_over_dataset` decorator helps enforce this as follows.

  ```py
  @seqio.map_over_dataset
  def translate(ex: Mapping[str, tf.Tensor],
                source_language: str,
                target_language: str) -> Mapping[str, tf.Tensor]:
    """Convert a translation dataset to a text2text pair.
<commit_msg>Clarify options for writing transformations with TensorFlow ops.<commit_after># SeqIO: Task-based datasets, preprocessing, and evaluation for sequence models.

**SeqIO** is a library for processing sequential data to be fed into downstream
sequence models. It uses [`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset)
to create scalable data pipelines but requires minimal use of TensorFlow. In
particular, with one line of code, the returned dataset can be transformed to a
numpy iterator and hence it is fully compatible with other frameworks such as
[JAX](https://github.com/google/jax) or
[PyTorch](https://pytorch.org/).

Currently, SeqIO assumes that the dataset is a sequence, i.e., each feature is
one-dimensional array. Modalities such as text or audio are naturally supported.
Images are supported as long as they are represented as sequences (e.g.,
[Image GPT](http://proceedings.mlr.press/v119/chen20s.html)). We will release this constraint
in the future in order to support higher dimensional data.

SeqIO is a refactor of the [`t5.data`](https://github.com/google-research/text-to-text-transfer-transformer/)
library used (in conjunction with the [Mesh Tensorflow](https://github.com/tensorflow/mesh)
Transformer implementation) to train the T5 models introduced in
[_Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer_](https://arxiv.org/abs/1910.10683).

If you have used `t5.data` in the past and want to know how SeqIO differs, please read [this section](#differences-from-t5data).

<!--Uncomment once forked.
## Installation

### From Pypi

```sh
pip install seqio
```

### From Source

```sh
git clone https://github.com/google/seqio.git
cd seqio
pip install -e .
-->

## Usage Tutorial

At a high level, we use SeqIO with the following steps.

  1. Define a `Task` (and optionally a `Mixture`).

  1. Define (or use an existing) a `FeatureConverter` based on the model architecture.

  1. Use the top-level function `seqio.get_dataset` to obtain the
     `tf.data.Dataset` instance.

We will look at each of these steps in detail.


### Defining a `Task`

The most important class in SeqIO is the `Task`. It is an abstraction that combines:

  * a raw *data source*
  * one or more *preprocessing* steps
  * a *vocabulary* to tokenize/detokenize each preprocessed feature for the model
  * a *postprocessor* to convert detokenized model outputs into a format for evaluation
  * one or more *metrics* to evaluate with

Oftentimes a `Task` lines up with a common benchmark. In this tutorial, we use
[WMT 19 English-German](http://www.statmt.org/wmt19/translation-task.html) machine
translation task. In the end, our `Task` will look like this:


```py
seqio.TaskRegistry.add(
    "wmt19_ende",
    seqio.TfdsDataSource(tfds_name="wmt19_translate/de-en:1.0.0"),
    preprocessors=[
        functools.partial(
            translate, source_language='en', target_language='de'),
        seqio.preprocessors.tokenize, seqio.preprocessors.append_eos
    ],
    output_features={
        'inputs':
            seqio.Feature(
                seqio.SentencePieceVocabulary('/path/to/inputs/vocab'),
                add_eos=False,
                dtype=tf.int32),
        'targets':
            seqio.Feature(
                seqio.SentencePieceVocabulary('/path/to/targets/vocab'),
                add_eos=True,
                dtype=tf.int32),
    },
    metric_fns=[bleu])
```

We typically add the `Task` to the global registry when we define it (as shown
above) to make it easier to use with model configs and flags. Thus, it  must
have a unique string name (`"wmt19_ende"` in this case). Note, however, that
you may also instantiate a `seqio.Task` directly without adding it to the
registry, if desired.

We'll now break down each part of the task definition.

#### Data Source

Data sources are the first step in your pipeline, providing a way to load raw
data in many formats as a `tf.data.Dataset`.
All data sources are subclasses of the `DataSource` base class and are defined in
[dataset_providers](https://github.com/google/seqio/tree/main/seqio/dataset_providers.py),

Existing implementations include:

  * `TfdsDataSource` for loading examples from [TensorFlow Datasets](https://www.tensorflow.org/datasets).
  * `TextLineDataSource` for loading examples from text files (e.g., tsv).
  * `TFExampleDataSource` for loading [`tf.train.Example`](https://www.tensorflow.org/tutorials/load_data/tfrecord) protos from a file (e.g. a `TFRecord` file.)
  * `FunctionDataSource` for providing an custom function that returns a `tf.data.Dataset`.

In our example, we are using the `TfdsDataSource`. We specify the name of the WMT dataset in TFDS ([`"wmt19_translate"`](https://www.tensorflow.org/datasets/catalog/wmt19_translate)), the specific config for the language pair that excludes the context for the open domain setting (`"de-en"`), and the version number (`"1.0.0"`).

#### Output Features

The `output_features` field expects a dictionary that maps string feature names
to `seqio.Feature` objects. This defines what the `Task` is expected to produce
in its output examples. The output examples *may* contain additional fields, but
they *must* contain these fields in the specified format or exceptions will be
raised.


Each `Feature` includes:

  * A `vocabulary`, which must subclass [`seqio.Vocabulary`](https://github.com/google/seqio/tree/main/seqio/vocabularies.py), to specify how the feature can be tokenized and detokenized. You may use `seqio.PassThroughVocabulary` if tokenization is not necessary.
  * `add_eos`, which specifies whether the feature should end with the vocabulary's EOS token.
  * The output `dtype` which must be a `tf.dtypes.DType`.

**Note:** specifying these options on `Feature` does not by itself ensure the proper transformations are applied -- you must also include the necessary preprocessors.

The [tasks used in T5](TODO) all produce "inputs" and "targets" features to be consumed by the text-to-text model. For a decoder-only language model, only a single feature (e.g., "targets") would be necessary.
Nevertheless, SeqIO is flexible enough to generate arbitrary output features what will be converted into model features by the [`FeatureConverter`](#featureconverter) later in the pipeline.

#### Preprocessors

Preprocessors are functions that transform one `tf.data.Dataset` into a new `tf.data.Dataset`. Typically this involves executing a `map` over the given dataset. The preprocessors provided to the `Task` will be executed sequentially.


As an example, let's look at the previously undefined `translate` from the "wmt19_ende" example above.

```py
def translate(dataset: tf.data.Dataset,
              source_language: str,
              target_language: str) -> tf.data.Dataset:
  def _translate(ex: Mapping[str, tf.Tensor]) -> Mapping[str, tf.Tensor]:
    """Convert a translation example to a text2text pair.

    For example, say the dataset returns examples of this format:
      {'de': 'Das ist gut.', 'en': 'That is good.'}
    If source_language = 'de', target_language = 'en', then the outputs will have
    the format:
      {'inputs': 'translate de to en: Das ist gut.',
      'targets': 'That is good.'}

    Args:
      x: an example to process.
      source_language: source language code (e.g. 'en') to translate from.
      target_language: target language code (e.g. 'de') to translate to.

    Returns:
      A preprocessed example with the format listed above.
    """
    src_str = f'translate {source_language}'
    tgt_str = f' to {target_language}: '
    return {
        'inputs': tf.strings.join([src_str, tgt_str, ex[source_language]]),
        'targets': ex[target_language],
    }

  return dataset.map(_translate,
                     num_parallel_calls=tf.data.experimental.AUTOTUNE)
```

The TFDS dataset provides the dataset where each example has the form: `{'de':
'Das ist gut.', 'en': 'That is good.'}`. We convert this to "inputs" and
"targets" with the appropriate prompt to inform the model of the task.


A few **important** notes:

  1. When instantiating a `Task`, the preprocessor functions can have the following arguments: `dataset`, `output_features`, and `sequence_length`. The first (positional) dataset argument is always required. If an argument named `output_features` is provided, the [output feature mapping](#output-features) will be passed to the preprocessor. If `sequence_length` is provided, a mapping from feature name to its *maximum* final sequence length ([provided by the caller](#getting-a-preprocessed-dataset) will be passed -- any sequences that are too long after preprocessing will be automatically truncated. If a preprocessor function does have other arguments, they must have default values or be bound (e.g., with `functools.partial` as used in `translate`) before instantiating the `Task`.

  1. Mapping functions operate on and return `tf.Tensor`s using TensorFlow operations. This is more flexible than it may sound:

    *  Automatic [AutoGraph](https://www.tensorflow.org/guide/function#autograph_transformations) conversion allow you to write python control flow in your transformations.
    * [tf.experimental.numpy](https://www.tensorflow.org/guide/tf_numpy) provides a numpy interface.
    * [`tf.py_function`](https://www.tensorflow.org/api_docs/python/tf/py_function) allows you to wrap arbitrary Python code. Note: `tf.data` pipelines using this function can only be run in the python process where they were defined, and performance is limited by the python GIL.

   See `tf.data.Dataset` [documentation](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) for more details.

  1. When calling `map`, it is important to **always** set `num_parallel_calls=tf.data.experimental.AUTOTUNE` to avoid creating a bottleneck. The `seqio.map_over_dataset` decorator helps enforce this as follows.

  ```py
  @seqio.map_over_dataset
  def translate(ex: Mapping[str, tf.Tensor],
                source_language: str,
                target_language: str) -> Mapping[str, tf.Tensor]:
    """Convert a translation dataset to a text2text pair.

{'fraction_non_alphanumeric': 0.09729946390099704, 'fraction_numerical': 0.005611503582343805, 'mean_word_length': 3.983770287141074, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 31, 'lorem ipsum': 0, 'www.': 16, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '273004', 'n_tokens_mistral': 5782, 'n_tokens_neox': 5548, 'n_words': 2280}