How to just use the mixture functionality in seqio

[StephennFernandes]

Hey there, I've been wanting to pretrain MT5 on Huggingface training script as mentioned here: https://github.com/huggingface/transformers/blob/main/examples/flax/language-modeling/run_t5_mlm_flax.py

But sadly the Huggingface script doesn't support a mixture to pretrain MT5 in such a way that the model generalise well on low-resource as well as high-resource langauges.

Hence I've been wanting to use the mixture functionality of seqio, but sadly upon using it i have to tokenize the model into the T5 sentencepiece vocabulary and seqio tasks does all the preprocessing.

The Huggingface trainer takes care of the preprocessing maping the dataset to the tokenizer etc.

My question is is there a way where i could only just use the mixture functionality of seqio without actually doing any preprocessing on the incoming datasets.

I was wondering if there is a way to feed in multiple datasets, get an output dataset (in text str format) which is only an appropriate mixture of all samples of the datsets, passed by the mixture function. which i could then use to pretrain on the HF trainer and then do all the preprocessing on it in HF trainer

[gauravmishra]

Hey Stephen, you can remove the preprocessors you don't need from your seqio Task definition. E.g. you can remove the seqio.preprocessors.tokenize preprocessor [1] to remove tokenization, and in your output_features[2] field pass a seqio.PassThroughVocabulary. See the linked documentation for more details.

[1] https://seqio.readthedocs.io/en/latest/overview.html#preprocessors
[2] https://seqio.readthedocs.io/en/latest/overview.html#output-features

[StephennFernandes]

hey @gauravmishra, thank you enormously for replying and helping out. I actually did as you instructed.

But sadly i get an error ValueError: Task dataset has incorrect rank for feature 'targets' after preprocessing: Got 0, expected 1

here's the following code for the same:

import functools

import seqio
import tensorflow as tf
import t5.data
from datasets import load_dataset
from t5.data import postprocessors
from t5.data import preprocessors
from t5.evaluation import metrics
from seqio import FunctionDataSource, utils

TaskRegistry = seqio.TaskRegistry



def gen_dataset(split, shuffle=False, seed=None, column="text", dataset_params=None):
    dataset = load_dataset(**dataset_params)
    if shuffle:
        if seed:
            dataset = dataset.shuffle(seed=seed)
        else:
            dataset = dataset.shuffle()
    while True:
        for item in dataset[str(split)]:
            yield item[column]


def dataset_fn(split, shuffle_files, seed=None, dataset_params=None):
    return tf.data.Dataset.from_generator(
        functools.partial(gen_dataset, split, shuffle_files, seed, dataset_params=dataset_params),
        output_signature=tf.TensorSpec(shape=(), dtype=tf.string, name=dataset_name)
    )


@utils.map_over_dataset
def target_to_key(x, key_map, target_key):
    """Assign the value from the dataset to target_key in key_map"""
    return {**key_map, target_key: x}



dataset_name = 'oscar-corpus/OSCAR-2109'
subset= 'mr'
dataset_params = {"path": dataset_name, "language":subset, "use_auth_token":True}
dataset_shapes = None

TaskRegistry.add(
    "oscar_marathi_corpus",
    source=seqio.FunctionDataSource(
        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
        splits=("train", "validation"),
        caching_permitted=False,
        num_input_examples=dataset_shapes,
    ),
    preprocessors=[
        functools.partial(
            target_to_key, key_map={
                "inputs": None,
                "targets": None,
            }, target_key="targets")],
    output_features={"targets": seqio.Feature(vocabulary=seqio.PassThroughVocabulary, add_eos=False, dtype=tf.string)},
    metric_fns=[]
)

dataset = seqio.get_mixture_or_task("oscar_marathi_corpus").get_dataset(
    sequence_length=None,
    split="train",
    shuffle=True,
    num_epochs=1,
    shard_info=seqio.ShardInfo(index=0, num_shards=10),
    use_cached=False,
    seed=42
)
for _, ex in zip(range(5), dataset.as_numpy_iterator()):
  print(ex)

ERROR:
ValueError: Task dataset has incorrect rank for feature 'targets' after preprocessing: Got 0, expected 1

i have mentioned tf.string in the task output_feature output_features={"targets": seqio.Feature(vocabulary=seqio.PassThroughVocabulary, add_eos=False, dtype=tf.string)}

Am i doing something wrong ?

i went through the docs to find a fix, but couldn't find anything.

[gauravmishra]

Hey, setting the rank of your feature to 0 should work:
output_features={"targets": seqio.Feature(vocabulary=seqio.PassThroughVocabulary, add_eos=False, dtype=tf.string, rank=0)}

[StephennFernandes]

@gauravmishra, thanks a ton once again. But sadly i am still facing issues here. I am really sorry to prompt you constantly.

after setting rank=0, and running:

for _, ex in zip(range(5), dataset.as_numpy_iterator()): print(ex)

i got a TypeError: ``tf.data.Dataset.as_numpy_iterator() is not supported for datasets that produce values of type <class 'tensorflow.python.data.util.structure.NoneTensor'>

so i tried to iter without calling the .as_numpy_iterator()

like this:

for _, ex in zip(range(5), dataset):
  print(ex)

And i get the output in byte representation :
{'inputs': None, 'targets': <tf.Tensor: shape=(), dtype=string, numpy=b'\xe0\xa4\x86.\xe0\xa4\xac\xe0\xa4\xbe\xe0\xa4\xb3\xe0\xa4\xbe\xe0\xa4\xb8\xe0\xa4\xbe\xe0\xa4\xb9\xe0\xa5\x87\xe0\xa4\xac \xe0\xa4\x86\xe0\xa4\x9c\xe0\xa4\xac\xe0\xa5\x87 \xe0\xa4\xaf\xe0\xa4\xbe\xe0\xa4\x82\xe0\xa4\x9a\xe0\xa4\xbe \xe0\xa4\x85\xe0\xa4\xad\xe0\xa5\x82\xe0\xa4\xa4\xe0\xa4\xaa\xe0\xa5\x82\xe0\xa4\xb0\xe0\xa5\x8d\xe0\xa4\xb5 \xe0\xa.'>}

where the item is a tensorflow.python.framework.ops.EagerTensor typed object, and thus cannot clearly decode bytes to string representation.

[gauravmishra]

Hey Stephen, you should remove the "inputs" key from your target_to_key preprocessor -

preprocessors=[
functools.partial(
target_to_key, key_map={
"targets": None,
}, target_key="targets")],

This should get rid of the "inputs" field that is getting set to None.

[StephennFernandes]

@gauravmishra, Hey Gaurav. i did as you instructed, but the output is still in bytes format:
and i could not convert bytes to utf-8 string, because the value is a tensorflow.python.framework.ops.EagerTensor

The following is the code for the same:

import functools

import seqio
import tensorflow as tf
import t5.data
from datasets import load_dataset
from t5.data import postprocessors
from t5.data import preprocessors
from t5.evaluation import metrics
from seqio import FunctionDataSource, utils

TaskRegistry = seqio.TaskRegistry



def gen_dataset(split, shuffle=False, seed=None, column="text", dataset_params=None):
    dataset = load_dataset(**dataset_params)
    if shuffle:
        if seed:
            dataset = dataset.shuffle(seed=seed)
        else:
            dataset = dataset.shuffle()
    while True:
        for item in dataset[str(split)]:
            yield item[column]


def dataset_fn(split, shuffle_files, seed=None, dataset_params=None):
    return tf.data.Dataset.from_generator(
        functools.partial(gen_dataset, split, shuffle_files, seed, dataset_params=dataset_params),
        output_signature=tf.TensorSpec(shape=(), dtype=tf.string, name=dataset_name)
    )


@utils.map_over_dataset
def target_to_key(x, key_map, target_key):
    """Assign the value from the dataset to target_key in key_map"""
    return {**key_map, target_key: x}



dataset_name = 'oscar-corpus/OSCAR-2109'
subset= 'mr'
dataset_params = {"path": dataset_name, "language":subset, "use_auth_token":True}
dataset_shapes = None

TaskRegistry.add(
    "oscar_marathi_corpus",
    source=seqio.FunctionDataSource(
        dataset_fn=functools.partial(dataset_fn, dataset_params=dataset_params),
        splits=("train", "validation"),
        caching_permitted=False,
        num_input_examples=dataset_shapes,
    ),
preprocessors=[
functools.partial(
target_to_key, key_map={
"targets": None,
}, target_key="targets")],
    output_features={"targets": seqio.Feature(vocabulary=seqio.PassThroughVocabulary, add_eos=False, dtype=tf.string, rank=0)},
    metric_fns=[]
)

dataset = seqio.get_mixture_or_task("oscar_marathi_corpus").get_dataset(
    sequence_length=None,
    split="train",
    shuffle=True,
    num_epochs=1,
    shard_info=seqio.ShardInfo(index=0, num_shards=10),
    use_cached=False,
    seed=42
)
for _, ex in zip(range(5), dataset):
  print(ex)
  break

output:
{'targets': <tf.Tensor: shape=(), dtype=string, numpy=b'\xe0\xa4\xb8\xe0\xa4\x82\xe0\xa4\xb0\xe0\xa4\x95\xe0\xa5\x8d\xe0\xa4\xb7\xe0\xa4\xa3 \xe0\xa4\xb8\xe0\xa4\x82\xe0\xa4\xb6\xe0\xa5\x8b\xe0\xa4\xa7\xe0\xa4\xa8 \xe0\xa4\xb5 \xe0\xa4\xb5\xe0\xa4\xbf\xe0\xa4\x95\xe0\xa4\xbe\xe0\xa4\xb8 \xe0\xa4\xb8\xe0\xa4\x82\xe0\xa4\x98\xe0\xa4\x9f\xe0\xa4\xa8\xe0\xa5\x87\xe0\xa4\xa4 \xe2\x80\x98\xe0\xa4\x85\xe0\xa4\xaa\xe0\xa5\x8d\xe0\xa4\xb0\xe0\xa5\x87\xe0\xa4\x82\xe0\xa4\x9f\xe0\xa4\xbf\xe0\xa4\xb8\xe2\x80\x99 \xe0\xa4\xaa\xe0\xa4\xa6\xe0\xa4\xbe\xe0\xa4\x9a\xe0\xa5\x8d\xe0\xa4\xaf\xe0\xa4\xbe 116 \xe0\xa4\x9c\xe0\xa4\xbe\xe0\xa4\x97\xe0\xa4\xbe\xe0\xa4\x82\xe0\xa4\xb8\xe0\xa4\xbe\xe0\xa4\xa0\xe0\xa5\x80 \xe0\xa4\xaa\xe0\xa4\xbe\xe0\xa4\xa4\xe0\xa5\x8d\xe0\xa4\xb0 \xe0\xa4\x89\xe0\xa4\xae\xe0\xa5\x87\xe0\xa4\xa6\xe0\xa4\xb5\xe0\xa4\xbe\xe0\xa4\xb0\xe0\xa4\xbe\xe0\xa4\x95\xe0\xa4\xa1\xe0\xa5\x82\xe0\xa4\xa8 \xe0\xa4\x85\xe0\xa4\xb0\xe0\xa5\x8d\xe0\xa4\x9c \xe0\xa4\xae\xe0\xa4\xbe\xe0\xa4\x97\xe0\xa4\xb5\xe0\xa4\xbf\xe0\xa4\xa3\xe0\xa5\x8d\xe0\xa4\xaf\xe0\xa4\xbe\xe0\xa4\xa4.'>}

[gauravmishra]

Something like the following should work -
for _, ex in zip(range(5), dataset):
print(ex['targets'].numpy().decode())
break

or if huggingface works with numpy, then you can pass dataset.as_numpy_iterator()

[StephennFernandes]

@gauravmishra Hey gaurav, It works fine! ... Cannot thank you enough for all the dedicated help. It really means a lot. Thanks a ton 🙏

[StephennFernandes]

@gauravmishra Hey Gaurav, I had a doubt. what is the best way to decide on which mixture ratio is optimal. In the mT5 paper the alpha value 0.3 gave the best balance between ideal performance for high and low resource languages.

However i am pretraining mT5 on indian languages, and i have a diverse variety of indian multi-lingual corpus, where Hindi has 60M+ samples and Kashmiri has around 100k samples.

So i wanted to know if i could h-param tune somehow on t5x, or would just using alpha=0.3 work fine in my usecase ?