HuggingFace datasets wrapper

[albertz]

We should be able to use HuggingFace datasets directly in RETURNN.

I guess the most canonical way would be to write a RETURNN Dataset for this. Maybe derived from CachedDataset2.

A separate independent more direct PyTorch dataset wrapper might make sense. Or actually I think not needed, as HuggingFace already directly supports this?

(cc @dthulke @NeoLegends @robin-p-schmitt)

[dthulke]

I guess the most canonical way would be to write a RETURNN Dataset for this. Maybe derived from CachedDataset2.

I already implemented this as a custom dataset some time ago (see implementation below).

Things to discuss are maybe how to handle tokenisation and other preprocessing steps.
In the current implementation there is either the option to define a map function (see https://huggingface.co/docs/datasets/about_map_batch) in the config file or to provide a preprocessed dataset stored with save_to_disk.

num_outputs and data_type is currently extracted from the dataset features attribute (which may not always be available).

Another point to discuss is how to handle caching for large datasets (currently the default caching mechanism of hf datasets is used which does not fit very well with our setups).

Code

import numpy
from returnn.datasets.basic import DatasetSeq
from returnn.datasets.cached2 import CachedDataset2
from returnn.util.basic import OptionalNotImplementedError


class HuggingfaceDataset(CachedDataset2):

  @staticmethod
  def kwargs_update_from_config(config, kwargs):
    super().kwargs_update_from_config(config, kwargs)
    if 'map_func' in kwargs:
      if isinstance(kwargs['map_func'], str):
        kwargs['map_func'] = config.typed_value(kwargs['map_func'])

  def __init__(self, dataset_opts, map_func=None, map_func_args=None, data_key='data', seq_tag_key='id', features=None,
               **kwargs):
    super(HuggingfaceDataset, self).__init__(**kwargs)

    self._seq_order = None

    self.dataset_opts = dataset_opts

    if isinstance(map_func, str):
      from returnn.config import get_global_config
      config = get_global_config(raise_exception=False)
      map_func = config.typed_value(map_func)

    if map_func_args is not None:
      map_func = map_func(**map_func_args)
    self.map_func = map_func
    self.dataset = None
    self.data_key = data_key
    self.seq_tag_key = seq_tag_key

    self.feature_keys = features

    self.data_dtype = {}

  def initialize(self):
    # Load the dataset
    import datasets
    if isinstance(self.dataset_opts, dict):
      self.dataset = datasets.load_dataset(**self.dataset_opts)
    else:
      self.dataset = datasets.load_from_disk(self.dataset_opts)
      assert isinstance(self.dataset, datasets.Dataset)
    if self.map_func is not None:
      self.dataset = self.map_func(self.dataset)
    if self.feature_keys is None:
      self.feature_keys = list(self.dataset.features.keys())
      if self.seq_tag_key is not None and self.seq_tag_key in self.feature_keys:
        self.feature_keys.remove(self.seq_tag_key)
      else:
        assert False, "Dataset does not have a seq_tag"

    self.dataset.set_format('numpy')

    if self.seq_tag_key is not None:
      assert self.seq_tag_key in self.dataset.column_names

    self.labels = {}
    self.num_outputs = {}
    for key in self.feature_keys:
      feature = self.dataset.features[key]
      dtype = None
      num_classes = None
      spatial_dims = 0
      while type(feature) is datasets.features.Sequence:
        spatial_dims += 1
        if feature.length != -1:
          num_classes = feature.length
        feature = feature.feature
      if type(feature) is datasets.features.ClassLabel:
        self.labels[key] = feature.names
        dtype = feature.dtype
        num_classes = feature.num_classes
      elif type(feature) is datasets.features.Value:
        dtype = feature.dtype
      elif isinstance(feature, (datasets.features.Array2D, datasets.features.Array3D, datasets.features.Array4D)):
        dtype = feature.dtype
        num_classes = feature.shape[-1]
        spatial_dims += len(feature.shape)
      else:
        assert False, f"Unsupported feature type {type(feature)}"

      len_shape = spatial_dims
      self.num_outputs[key] = [num_classes, len_shape]

      self.data_dtype[key] = dtype

    super().initialize()

  def get_data_dim(self, key):
    if key in self.num_outputs:
      return self.num_outputs[key][0]
    return super().get_data_dim(key)

  def get_data_dtype(self, key):
    return self.data_dtype[key]

  def _get_seq_len(self, seq_idx):
    return len(self.dataset[seq_idx][self.data_key])

  @property
  def num_seqs(self):
    assert self._seq_order is not None, "num_seqs is only known after calling init_seq_order()"
    return len(self._seq_order)

  def get_tag(self, sorted_seq_idx):
    return self.dataset[int(self.get_corpus_seq_idx(sorted_seq_idx))][self.seq_tag_key]

  def get_all_tags(self):
    return list(self.dataset[self.seq_tag_key])

  def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
    """
    :param int|None epoch:
    :param list[str]|None seq_list: List of sequence tags, to set a predefined order.
    :param list[int]|None seq_order: List of corpus sequence indices, to set a predefined order.
    :rtype: bool
    :returns whether the order changed (True is always safe to return)
    """
    super(HuggingfaceDataset, self).init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)

    if seq_order:
      self._seq_order = seq_order
      # TODO can we return False?
      return True

    if seq_list:
      all_tags = self.get_all_tags()
      self._seq_order = [all_tags.index(tag) for tag in seq_list]
      # TODO can we return False?
      return True

    try:
      self._seq_order = self.get_seq_order_for_epoch(
        epoch=epoch, num_seqs=self.dataset.num_rows, get_seq_len=self._get_seq_len)
    except OptionalNotImplementedError:
      # only support seq_ordering that need no length here
      assert self.seq_ordering in ["default", "reverse", "random"]
      self._seq_order = self.get_seq_order_for_epoch(
        epoch=epoch, num_seqs=self.dataset.num_rows, get_seq_len=None)

    return True

  def _collect_single_seq(self, seq_idx):
    """
    :param int seq_idx: sorted seq idx
    :return:
    """
    corpus_seq_idx = self.get_corpus_seq_idx(seq_idx)

    def ensure_numpy(x):
      if not isinstance(x, numpy.ndarray):
        return numpy.array(x)
      return x

    dataset_item = self.dataset[int(corpus_seq_idx)]
    features = {f: ensure_numpy(dataset_item[f]) for f in self.feature_keys}
    return DatasetSeq(
      seq_idx,
      features=features,
      targets=None,
      seq_tag=dataset_item[self.seq_tag_key]
    )

  def get_current_seq_order(self):
    """
    :rtype: list[int]
    """
    assert self._seq_order is not None
    return self._seq_order

  def get_corpus_seq_idx(self, sorted_seq_idx):
    """
    :param int sorted_seq_idx:
    :return corpus_seq_idx
    :rtype: int
    """
    return self._seq_order[sorted_seq_idx]

  def can_serialize_data(self, key):
    return True

  def serialize_data(self, key, data):
    if key in self.labels:
      return super().serialize_data(key, data)
    if isinstance(data, numpy.ndarray):
      data = data.tolist()
    return data

[albertz]

Related is the Sisyphus job to prepare HuggingFace datasets (rwth-i6/i6_core#253). Doesn't this handle the caching? Ideally we should prepare our dataset wrapper here such that it works properly together with this download preparation job.

[albertz]

@dthulke Can you say some examples what HF datasets you use?

[dthulke]

Related is the Sisyphus job to prepare HuggingFace datasets (rwth-i6/i6_core#253). Doesn't this handle the caching? Ideally we should prepare our dataset wrapper here such that it works properly together with this download preparation job.

Yes, this handles the caching of the initial dataset download, but not the caching of the processed version (via dataset.map). But we could add a separate job for this.

@dthulke Can you say some examples what HF datasets you use?

In RETURNN, I mainly use hf datasets for sequence classification (e.g. sentiment analysis) or sequence tagging task (named entity recognition). For example: https://huggingface.co/datasets/conll2003

In addition, I have a few datasets the I load with custom dataset loading scripts or the default json dataset implementation.

One example, for additional preprocessing (beyond tokenisation) is to include document-level/cross-sentence context or to convert NER labels (given as start end positions) to BIO labels.

[albertz]

Is the information here still accurate/relevant/up-to-date?

What about post-processing, save_to_disk? Do you do that offline? Or is post-processing (map) also feasible on-the-fly?

I wonder also: If we would want to combine multiple HF datasets (e.g. YODAS + Librispeech or so), what would be the best way? E.g. via HF datasets concatenate_datasets or interleave_datasets? Is this how you would do it? This supports custom data ratios? You would then dump the result again with save_to_disk first? Or use some of the RETURNN ways? (E.g. ConcatDataset or DistributeFilesDataset or so. Also see #1700 for MixingDataset.)

[albertz]

I see that for post-processing (also potentially including data mixing), we can then in the end use HF datasets shard to store individual parts (or maybe the data ratios / mixing are not needed at that point), and then use DistributeFilesDataset over all the generated data files (and there we can control data ratios by just duplicating some entries).

[albertz]

Small comments on the implementation by @dthulke (here):

We maybe should not do the load_dataset in initialize. This will first load it in the main proc, but then also additionally in each data loader worker. Rather, we should maybe make this lazily, do the load_dataset on first usage, as many other datasets are also doing it (see e.g. OggZipDataset, LmDataset, etc).

We maybe should let the user define an explicit list/mapping of the features, and maybe also their shape/dtype? (This might also be a requirement if we do lazy loading. Not sure...)

get_global_config should not be used. (I guess this was for serialization issues? We now have several other better ways to do it. Or with i6_experiments.users.zeyer.serialization_v2.serialize_config, there is no need to take any extra care.)

get_all_tags is likely slow. This should be cached on first usage.

I'm not sure why you are catching OptionalNotImplementedError for the get_seq_order_for_epoch. I don't see how this can be thrown there? The _get_seq_len that you implemented should always work?

The data_key is only used for _get_seq_len, for nothing else. Is this intended that way? Then maybe it should be renamed to seq_len_feature_key or so? Also, is "data" a reasonable default? Is this often used in HF datasets? At least for the speech datasets, I never see "data". Instead, there is "text", "transcription", "audio", etc. I sometimes also see explicit length info like "audio_length_s", "duration".

You do random access here (self.dataset[int(corpus_seq_idx)] (btw, why the int cast here?)). Is this efficient?

I guess it anyway makes sense to wrap this inside a MultiProcDataset or DistributeFilesDataset to use background prefetching?

[albertz]

You do random access here (self.dataset[int(corpus_seq_idx)] (btw, why the int cast here?)). Is this efficient?

On the speed: In the HF datasets process doc on shuffling (shuffle):

Shuffling takes the list of indices [0:len(my_dataset)] and shuffles it to create an indices mapping. However as soon as your Dataset has an indices mapping, the speed can become 10x slower. This is because there is an extra step to get the row index to read using the indices mapping, and most importantly, you aren’t reading contiguous chunks of data anymore.

shuffle sounds very much like what we do in init_seq_order, i.e. building a list of indices (indices mapping). Then (if it is not the default order), we aren’t reading contiguous chunks of data anymore.

So, is this an issue? This sounds like an issue? Is it still an issue when we have the data on local disk? (Does the format matter? What about Parquet? Arrow? save_to_disk uses Arrow?)

The doc also mentions "fast approximate shuffling IterableDataset.shuffle()". As I understand the implementation, it prefetches some amount of entries into a buffer, but this is done from the beginning of the dataset. Then it randomly samples an item from the buffer. I see that this is fast now, because we again can read contiguous chunks of data. But at the same time, this is not really so random at all? This sounds quite bad? However, it seems this would also at least shuffle over shards, if there are shards. That makes it better.

We could maybe also adapt our shuffling to operate more in chunks. I.e. first segment the data into chunks of some fixed chunk size, then shuffle the chunks, then shuffle within each chunk. That way, most of the random access would be more close together. But I don't know if the underlying Arrow dataset would already benefit from that, or still needs extra logic (like prefetching the whole chunk). Also, not sure about a good chunk size.

I read a bit more about that. E.g. here on HF datasets arrow:

Arrow’s standard format allows zero-copy reads which removes virtually all serialization overhead.
You can obtain the best performance by accessing slices of data (or “batches”), in order to reduce the amount of lookups on disk.

They have this example code:

wiki = load_dataset("wikipedia", "20220301.en", split="train")
batch_size = 1000
for i in range(0, len(wiki), batch_size):
    batch = wiki[i:i + batch_size]

So, it seems that the underlying Arrow would not internally do this, and I need to take care of such logic (?).

But that wouldn't really be such a big issue.

Edit I just checked datasets.arrow_dataset.Dataset.__iter__ (that is what a common dataset __iter__ will likely be):

    def __iter__(self):
        if self._indices is None:
            # Fast iteration
            # Benchmark: https://gist.github.com/mariosasko/0248288a2e3a7556873969717c1fe52b (fast_iter_batch)
            ...
            batch_size = config.ARROW_READER_BATCH_SIZE_IN_DATASET_ITER
            for pa_subtable in table_iter(self.data, batch_size=batch_size):
                for i in range(pa_subtable.num_rows):
                    ...
                    yield formatted_output
        else:
            for i in range(self.num_rows):
                yield self._getitem(i)

So you see why it is slower with shuffle (where self._indices would be set): It doesn't use the batched reading.

However, if we made sure that even with shuffling, most of the seq indices are still closeby, I think we still can use batched reading, and speeding it up again quite a bit.

-> I opened a separate issue #1777 to discuss these potential speed optimizations.

[albertz]

how to handle tokenisation

There could be some extra keys, defined maybe like: tokenize_keys={"labels": ("text", {...vocab_opts...}), ...}. The vocab opts would be for Vocabulary.create_vocab, and it would then use Vocabulary.get_seq to map from "text" (asumed to be string) to labels, under the data key "labels".

(PostprocessingDataset could also handle that, but having this directly here makes it slightly more convenient. This is probably also a very common thing you want to use.)

[albertz]

One other small aspect: I just learned: When you have used save_to_disk on the dataset, you cannot load that via load_dataset (huggingface/datasets#7018). You must use DatasetDict.load_from_disk or Dataset.load_from_disk. This is quite unfortunate, as this is a valid case that we might have (e.g. rwth-i6/i6_core#627). So it means, the dataset in RETURNN should somehow support both (or all three) cases: load_dataset, DatasetDict.load_from_disk, Dataset.load_from_disk. And we must treat these cases differently.

[albertz]

Hm, another issue: Loading a set of arrow files directly, that's not really supported via Dataset.load_from_disk (which can load only a single file), so for that use case (which is a valid use case that we will have when using DistributeFilesDataset with it), again we need to see how to do that. I think for this case, somehow load_dataset can work, but not exactly sure how. Or otherwise we need to see how to make it work (but it should be possible somehow).