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?

[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

```python 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 (https://github.com/rwth-i6/i6_core/pull/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 (https://github.com/rwth-i6/i6_core/pull/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.