returnn 1.20251006.114241__py3-none-any.whl → 1.20251007.223754__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251006.114241
+Version: 1.20251007.223754
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn/_setup_info_generated.py

@@ -1,2 +1,2 @@
-version = '1.20251006.114241'
-long_version = '1.20251006.114241+git.7745ba7'
+version = '1.20251007.223754'
+long_version = '1.20251007.223754+git.eb1103a'

returnn/datasets/basic.py

@@ -19,6 +19,7 @@ import os
 import math
 import numpy
 import functools
+import types
 from typing import TYPE_CHECKING, Optional, Any, Set, Tuple, Union, Type, Dict, Sequence, List, Callable
 
 from returnn.log import log
@@ -154,7 +155,7 @@ class Dataset:
         self.seq_tags_filter = set(self._load_seq_list_file(seq_list_filter_file)) if seq_list_filter_file else None
         self.unique_seq_tags = unique_seq_tags
         self._seq_order_seq_lens_file = seq_order_seq_lens_file
-        self._seq_order_seq_lens_by_idx = None
+        self._seq_order_seq_lens_by_idx: Optional[Sequence[Union[int, float]]] = None
         # There is probably no use case for combining the two, so avoid potential misconfiguration.
         assert self.partition_epoch == 1 or self.repeat_epoch == 1, (
             "Combining partition_epoch and repeat_epoch is prohibited."
@@ -486,12 +487,8 @@ class Dataset:
         """
         raise NotImplementedError
 
-    def _get_seq_order_seq_lens_by_idx(self, seq_idx):
-        """
-        :param int seq_idx:
-        :rtype: int
-        """
-        if not self._seq_order_seq_lens_by_idx:
+    def _get_seq_order_seq_lens_by_idx(self, seq_idx: int) -> Union[int, float]:
+        if self._seq_order_seq_lens_by_idx is None:
             assert self._seq_order_seq_lens_file
             if self._seq_order_seq_lens_file.endswith(".gz"):
                 import gzip
@@ -502,11 +499,12 @@ class Dataset:
             seq_lens = eval(raw)
             assert isinstance(seq_lens, dict)
             all_tags = self.get_all_tags()
-            self._seq_order_seq_lens_by_idx = [seq_lens[tag] for tag in all_tags]
+            self._seq_order_seq_lens_by_idx = numpy.array([seq_lens[tag] for tag in all_tags])
+        self._get_seq_order_seq_lens_by_idx = self._seq_order_seq_lens_by_idx.__getitem__  # faster
         return self._seq_order_seq_lens_by_idx[seq_idx]
 
     def get_seq_order_for_epoch(
-        self, epoch: Optional[int], num_seqs: int, get_seq_len: Optional[Callable[[int], int]] = None
+        self, epoch: Optional[int], num_seqs: int, get_seq_len: Optional[Callable[[int], Union[int, float]]] = None
     ) -> Sequence[int]:
         """
         Returns the order of the given epoch.
@@ -515,7 +513,7 @@ class Dataset:
 
         :param epoch: for 'random', this determines the random seed
         :param num_seqs:
-        :param get_seq_len: function (originalSeqIdx: int) -> int
+        :param get_seq_len: function (originalSeqIdx: int) -> int|float
         :return: the order for the given epoch. such that seq_idx -> underlying idx
         """
         if epoch is None:
@@ -561,8 +559,9 @@ class Dataset:
             seq_index = range(num_seqs - 1, -1, -1)  # type: Union[range, Sequence[int]]
         elif seq_ordering_method in ["sorted", "sorted_reverse"]:
             assert get_seq_len
-            reverse = -1 if seq_ordering_method == "sorted_reverse" else 1
-            seq_lens = [reverse * get_seq_len(i) for i in range(num_seqs)]
+            seq_lens = _get_seq_len_as_array(get_seq_len, num_seqs)
+            if seq_ordering_method == "sorted_reverse":
+                seq_lens = -seq_lens
             seq_index = numpy.argsort(seq_lens, kind="stable")
         elif seq_ordering_method == "random" or seq_ordering_method.startswith("random:"):
             tmp = seq_ordering_method.split(":", 1)
@@ -628,7 +627,7 @@ class Dataset:
                 nth = 1
             else:
                 nth = int(tmp[1])
-            seq_lens = numpy.array([get_seq_len(i) for i in range(num_seqs)])
+            seq_lens = _get_seq_len_as_array(get_seq_len, num_seqs)
             rnd_seed = self._get_random_seed_for_epoch(epoch=epoch, num_epochs_fixed=nth)
             random_generator = numpy.random.RandomState(rnd_seed)
             seq_index = random_generator.permutation(num_seqs)  # type: Union[numpy.ndarray, List[int]]
@@ -1501,6 +1500,7 @@ def get_dataset_class(name: Union[str, Type[Dataset]]) -> Optional[Type[Dataset]
         "distrib_files",
         "postprocessing",
         "text_dict",
+        "huggingface",
     ]
     for mod_name in mod_names:
         mod = import_module("returnn.datasets.%s" % mod_name)
@@ -1757,3 +1757,19 @@ def set_config_extern_data_from_dataset(config, dataset):
         "extern_data",
         {key: _data_kwargs_from_dataset_key(dataset=dataset, key=key) for key in dataset.get_data_keys()},
     )
+
+
+def _get_seq_len_as_array(get_seq_len: Callable[[int], Union[int, float]], num_seqs: int) -> numpy.ndarray:
+    if num_seqs == 0:
+        return numpy.zeros((0,), dtype=numpy.int32)
+    if isinstance(get_seq_len, (types.BuiltinMethodType, types.MethodWrapperType, types.MethodType)):
+        # Call it once. This might trigger some caching.
+        get_seq_len(0)
+        # Get it again. This might now get us a different (cached) function, e.g. array.__getitem__.
+        get_seq_len = getattr(get_seq_len.__self__, get_seq_len.__name__)
+        assert isinstance(get_seq_len, (types.BuiltinMethodType, types.MethodWrapperType, types.MethodType))
+        obj = get_seq_len.__self__
+        if isinstance(obj, numpy.ndarray) and get_seq_len.__name__ == "__getitem__":
+            assert obj.shape == (num_seqs,)
+            return obj
+    return numpy.array([get_seq_len(i) for i in range(num_seqs)])

returnn/datasets/distrib_files.py

@@ -135,7 +135,7 @@ class DistributeFilesDataset(CachedDataset2):
     def __init__(
         self,
         *,
-        files: Union[List[FileTree], os.PathLike],
+        files: Union[List[FileTree], os.PathLike, Callable[[], List[FileTree]]],
         get_sub_epoch_dataset: Callable[[List[FileTree]], Dict[str, Any]],
         preload_next_n_sub_epochs: int = 1,
         buffer_size: int = 1,
@@ -151,6 +151,7 @@ class DistributeFilesDataset(CachedDataset2):
             can also be specified as a path to a .txt file containing one file per line,
             or a python file containing the repr of a list of arbitrarily nested python objects,
             or a JSON file containing a list of arbitarily nested (JSON) objects.
+            It can also be a callable which returns such a list.
         :param get_sub_epoch_dataset: callable which returns a dataset dict for a given subset of files
         :param preload_next_n_sub_epochs: how many sub epoch datasets to preload
         :param buffer_size: buffer size for each worker, number of seqs to prefetch
@@ -244,6 +245,11 @@ class DistributeFilesDataset(CachedDataset2):
             return
         if isinstance(self.files, list):
             self._files = self.files
+        elif callable(self.files):
+            self._files = self.files()
+            assert isinstance(self._files, list), (
+                f"{self}: callable files {self.files} must return a list, got {type(self._files)}"
+            )
         elif isinstance(self.files, (str, os.PathLike)):
             _, ext = os.path.splitext(self.files)
             assert ext, f"{self}: no file extension on file list file {self.files}"

returnn/datasets/huggingface.py

@@ -0,0 +1,434 @@
+"""
+HuggingFace dataset wrapper
+
+See https://github.com/rwth-i6/returnn/issues/1257 for some initial discussion.
+"""
+
+from __future__ import annotations
+from typing import TYPE_CHECKING, Optional, Union, Any, Callable, Sequence, Dict, List
+import os
+import re
+import numpy
+from returnn.tensor import Tensor
+from returnn.util import file_cache
+from .basic import DatasetSeq
+from .cached2 import CachedDataset2
+from .util.vocabulary import Vocabulary
+from .util.strings import str_to_numpy_array
+
+if TYPE_CHECKING:
+    # noinspection PyUnresolvedReferences,PyPackageRequirements
+    import datasets
+
+
+class HuggingFaceDataset(CachedDataset2):
+    """
+    HuggingFace dataset wrapper.
+    """
+
+    def __init__(
+        self,
+        dataset_opts: Union[
+            Dict[str, Any],
+            str,
+            os.PathLike,
+            Sequence[Union[str, os.PathLike]],
+            Callable[[], Union[Dict[str, Any], str, os.PathLike, Sequence[Union[str, os.PathLike]], datasets.Dataset]],
+        ],
+        *,
+        use_file_cache: bool = False,
+        map_func: Optional[Callable[[datasets.Dataset], datasets.Dataset]] = None,
+        rename_columns: Optional[Dict[str, str]] = None,
+        cast_columns: Optional[Dict[str, Dict[str, Any]]] = None,
+        data_format: Dict[str, Dict[str, Any]],
+        seq_tag_column: Optional[str] = "id",
+        sorting_seq_len_column_data: Optional[str] = None,
+        sorting_seq_len_column: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        :param dataset_opts: either a dict of options for :func:`datasets.load_dataset`
+            or a path to a local dataset for :func:`datasets.load_from_disk`,
+            or a list of Arrow filenames to load with :func:`datasets.Dataset.from_file` and concatenate.
+            It can also be a callable returning one of the above,
+            or returning a :class:`datasets.Dataset` directly.
+        :param use_file_cache: if True, will cache the dataset files on local disk using :mod:`file_cache`.
+            This only works for dataset_opts which is a str or list of str (or callable returning that).
+        :param map_func: optional function to apply to the dataset after loading
+        :param rename_columns: if given, will rename these columns
+        :param cast_columns: if given, will cast these columns to the specified types.
+            This is useful if the dataset has not the expected types.
+            See :func:`datasets.Dataset.cast` for details.
+            You can also e.g. enforce some sample_rate for audio, etc.
+        :param data_format:
+            For each column name (data key), specify the format,
+            as a dict with entries for "dim", "ndim", "shape", and/or "dtype",
+            compatible to :class:`Tensor`.
+            It can be a subset of the available columns.
+            If "vocab" is specified, and the underlying HF datasets column is of dtype "string",
+            it will automatically tokenize the string using the vocab.
+        :param seq_tag_column: key (column name) in the dataset to use as sequence tag.
+            If None, will use the sequence index as tag.
+        :param sorting_seq_len_column_data: key (column name) in the dataset to use for sorting by sequence length.
+            It will take len(dataset[sorting_seq_len_column_data]) as sequence length (only for sorting/shuffling).
+        :param sorting_seq_len_column: key (column name) in the dataset to use for sorting by sequence length.
+            It will take the value of dataset[sorting_seq_len_column] as sequence length (only for sorting/shuffling).
+            E.g. some datasets provide "duration", "duration_ms", "wav_filesize" or similar such information
+            which can be used.
+        """
+        super().__init__(**kwargs)
+
+        self.dataset_opts = dataset_opts
+        self.use_file_cache = use_file_cache
+        self.map_func = map_func
+        self.rename_columns = rename_columns
+        self.cast_columns = cast_columns
+
+        self.data_format: Dict[str, Tensor] = {k: _make_tensor_template(v, k) for k, v in data_format.items()}
+        self.seq_tag_column: Optional[str] = seq_tag_column
+        self.sorting_seq_len_column_data = sorting_seq_len_column_data
+        self.sorting_seq_len_column = sorting_seq_len_column
+
+        self.labels = {k: data.vocab.labels for k, data in self.data_format.items() if data.vocab}
+        self.num_outputs = {k: (data.dim, data.ndim) for k, data in self.data_format.items()}
+
+        self.hf_dataset: Optional[datasets.Dataset] = None  # lazily loaded, _lazy_init
+        self._seq_order: Optional[Sequence[int]] = None  # init_seq_order
+        self._seq_tags: Optional[List[str]] = None  # get_all_tags cache
+
+    def _lazy_init(self):
+        if self.hf_dataset is not None:
+            return
+
+        # Load the dataset
+        # noinspection PyUnresolvedReferences,PyPackageRequirements
+        import datasets
+
+        dataset_opts = self.dataset_opts
+        if callable(dataset_opts):
+            dataset_opts = dataset_opts()
+        if self.use_file_cache:
+            assert isinstance(dataset_opts, (str, os.PathLike, list, tuple)), (
+                f"{self}: with use_file_cache, dataset_opts must be str or list of str, got {type(dataset_opts)}"
+            )
+            if isinstance(dataset_opts, (str, os.PathLike)):
+                dataset_opts = get_arrow_shard_files_from_hf_dataset_dir(dataset_opts)
+            assert isinstance(dataset_opts, (list, tuple))
+            cache = file_cache.get_instance()
+            dataset_opts = [cache.get_file(os.fspath(fn)) for fn in dataset_opts]
+            self.set_file_cache(cache)
+        if isinstance(dataset_opts, dict):
+            self.hf_dataset = datasets.load_dataset(**dataset_opts)
+        elif isinstance(dataset_opts, str):
+            self.hf_dataset = datasets.load_from_disk(dataset_opts)
+        elif isinstance(dataset_opts, (list, tuple)):
+            self.hf_dataset = datasets.concatenate_datasets([datasets.Dataset.from_file(fn) for fn in dataset_opts])
+        elif isinstance(dataset_opts, datasets.Dataset):
+            self.hf_dataset = dataset_opts
+        else:
+            raise TypeError(f"{self}: invalid dataset_opts type {type(dataset_opts)}")
+        assert isinstance(self.hf_dataset, datasets.Dataset), (
+            f"{self}: Expected single dataset, got {type(self.hf_dataset)} {self.hf_dataset}. Specify split if needed."
+        )
+
+        if self.map_func is not None:
+            self.hf_dataset = self.map_func(self.hf_dataset)
+
+        if self.rename_columns:
+            self.hf_dataset = self.hf_dataset.rename_columns(self.rename_columns)
+
+        if self.cast_columns:
+            # Note: prefer cast_column, as this can avoid using `map`, i.e. be faster.
+            for key, column_format in self.cast_columns.items():
+                assert key in self.hf_dataset.features, (
+                    f"{self}: cast_column {key} not in dataset features {self.hf_dataset.features}"
+                )
+                feat = datasets.features.features.generate_from_dict(column_format)
+                self.hf_dataset = self.hf_dataset.cast_column(key, feat)
+
+        if self.seq_tag_column:
+            assert self.seq_tag_column in self.hf_dataset.features, (
+                f"{self}: seq_tag_column {self.seq_tag_column} not in dataset features {self.hf_dataset.features}"
+            )
+            assert self.hf_dataset.features[self.seq_tag_column].dtype in ("string", "int64"), (
+                f"{self}: seq_tag_column {self.seq_tag_column} must be of dtype string or int64,"
+                f" got {self.hf_dataset.features[self.seq_tag_column].dtype}"
+            )
+
+        selected_columns = list(self.data_format.keys())
+        if self.seq_tag_column and self.seq_tag_column not in selected_columns:
+            selected_columns.append(self.seq_tag_column)
+        if self.sorting_seq_len_column and self.sorting_seq_len_column not in selected_columns:
+            selected_columns.append(self.sorting_seq_len_column)
+        if self.sorting_seq_len_column_data and self.sorting_seq_len_column_data not in selected_columns:
+            selected_columns.append(self.sorting_seq_len_column_data)
+        self.hf_dataset = self.hf_dataset.select_columns(selected_columns)
+
+        self.hf_dataset.set_format("numpy")
+
+        for key, user_format in self.data_format.items():
+            feature = self.hf_dataset.features[key]
+            inferred_format = _infer_data_format_for_feature(feature, f"{self}: column {key}: ")
+            if user_format.vocab and inferred_format["dtype"] == "string":
+                pass  # allow to auto-tokenize strings when vocab is specified
+            else:
+                for key_ in ["dtype", "ndim", "dim"]:
+                    assert getattr(user_format, key_) == inferred_format[key_], (
+                        f"{self}: column {key}, user-specified {user_format}, {key_}:"
+                        f" user-specified {getattr(user_format, key_)} does not match inferred {inferred_format[key_]}"
+                    )
+            if "vocab" in inferred_format and not user_format.vocab:
+                assert user_format.sparse, f"{self}: column {key}: user_format expected to be sparse, got {user_format}"
+                user_format.sparse_dim.vocab = Vocabulary.create_vocab(**inferred_format["vocab"])
+                self.labels[key] = user_format.vocab.labels
+
+    def get_data_keys(self) -> List[str]:
+        """:return: list of data keys"""
+        return list(self.data_format.keys())
+
+    def get_target_list(self) -> List[str]:
+        """:return: list of target keys"""
+        return self.get_data_keys()  # it's somewhat arbitrary...
+
+    def get_data_shape(self, key: str) -> List[int]:
+        """:return: data shape for the given key"""
+        return list(self.data_format[key].shape)
+
+    def get_data_dim(self, key: str) -> int:
+        """:return: data dimension for the given key"""
+        return self.data_format[key].dim
+
+    def is_data_sparse(self, key: str) -> bool:
+        """:return: whether the data is sparse for the given key"""
+        return self.data_format[key].sparse
+
+    def get_data_dtype(self, key: str) -> str:
+        """:return: dtype"""
+        return self.data_format[key].dtype
+
+    def _get_seq_len(self, seq_idx: int) -> Union[int, float]:
+        if self._seq_order_seq_lens_by_idx is not None:
+            self._get_seq_len = self._seq_order_seq_lens_by_idx.__getitem__  # faster
+            return self._seq_order_seq_lens_by_idx[seq_idx]
+        assert not self._seq_order_seq_lens_file  # not expected to call this
+        if self.sorting_seq_len_column:
+            self._seq_order_seq_lens_by_idx = numpy.array(self.hf_dataset[self.sorting_seq_len_column])
+            self._get_seq_len = self._seq_order_seq_lens_by_idx.__getitem__  # faster
+            v = self._seq_order_seq_lens_by_idx[seq_idx]
+            return int(v)  # noqa
+        if self.sorting_seq_len_column_data:
+            v = self.hf_dataset[seq_idx][self.sorting_seq_len_column_data]
+            return len(v)  # noqa
+        raise ValueError(
+            f"{self}: sorting/shuffling by seq len not configured,"
+            f" need sorting_seq_len_column or sorting_seq_len_column_data"
+        )
+
+    @property
+    def num_seqs(self) -> int:
+        """:return: number of sequences"""
+        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: int) -> str:
+        """:return: tag of the sequence"""
+        corpus_seq_idx = self.get_corpus_seq_idx(sorted_seq_idx)
+        self._lazy_init()
+        dataset_item = self.hf_dataset[corpus_seq_idx]
+        return self._get_seq_tag(corpus_seq_idx, dataset_item)
+
+    def get_all_tags(self) -> List[str]:
+        """:return: all tags"""
+        if self._seq_tags is not None:
+            return self._seq_tags
+        self._lazy_init()
+        if self.seq_tag_column:
+            res = list(map(str, self.hf_dataset[self.seq_tag_column]))
+        else:
+            res = [f"seq-{i}" for i in range(self.hf_dataset.num_rows)]
+        self._seq_tags = res
+        return res
+
+    def get_total_num_seqs(self, *, fast: bool = False) -> int:
+        """:return: total number of sequences in the dataset"""
+        if fast:
+            return super().get_total_num_seqs(fast=True)
+        self._lazy_init()
+        return self.hf_dataset.num_rows
+
+    def init_seq_order(
+        self,
+        epoch: Optional[int] = None,
+        seq_list: Optional[Sequence[str]] = None,
+        seq_order: Optional[Sequence[int]] = None,
+    ) -> bool:
+        """
+        :param epoch:
+        :param seq_list: List of sequence tags, to set a predefined order.
+        :param seq_order: List of corpus sequence indices, to set a predefined order.
+        :returns whether the order changed (True is always safe to return)
+        """
+        super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
+
+        if seq_order is not None:
+            self._seq_order = seq_order
+        elif seq_list is not None:
+            all_tags = self.get_all_tags()
+            self._seq_order = [all_tags.index(tag) for tag in seq_list]
+        elif epoch is None:
+            self._seq_order = ()
+        else:
+            self._lazy_init()
+            self._seq_order = self.get_seq_order_for_epoch(
+                epoch=epoch, num_seqs=self.hf_dataset.num_rows, get_seq_len=self._get_seq_len
+            )
+        return True
+
+    def _collect_single_seq(self, seq_idx: int) -> DatasetSeq:
+        # noinspection PyUnresolvedReferences,PyPackageRequirements
+        import datasets
+
+        corpus_seq_idx = self.get_corpus_seq_idx(seq_idx)
+
+        def _ensure_numpy(k, x):
+            if isinstance(x, numpy.ndarray):  # fast path
+                return x
+            if isinstance(x, str):
+                if self.data_format[k].dtype == "string":
+                    return str_to_numpy_array(x)
+                if self.data_format[k].vocab:
+                    return numpy.array(self.data_format[k].vocab.get_seq(x), dtype=self.data_format[k].dtype)
+                raise ValueError(f"{self}: column {k}: cannot convert string {x!r} to numpy array")
+            feat = self.hf_dataset.features[k]
+            if isinstance(feat, datasets.features.Audio):
+                # In HF datasets 3, this is just a dict.
+                # In HF datasets 4, this can also be a datasets.features._torchcodec.AudioDecoder.
+                assert isinstance(x, dict) or x.__class__.__name__ == "AudioDecoder"
+                if feat.decode:
+                    x = x["array"]
+                else:
+                    x = x["bytes"]
+            if isinstance(x, numpy.ndarray):  # fast path
+                return x
+            if isinstance(x, (bytes, bytearray)):
+                return numpy.frombuffer(x, dtype=self.data_format[k].dtype)
+            return numpy.array(x)
+
+        self._lazy_init()
+        dataset_item = self.hf_dataset[corpus_seq_idx]
+        seq_tag = self._get_seq_tag(corpus_seq_idx, dataset_item)
+        features = {k: _ensure_numpy(k, dataset_item[k]) for k in self.data_format}
+        return DatasetSeq(seq_idx, features=features, seq_tag=seq_tag)
+
+    def _get_seq_tag(self, corpus_seq_idx: int, dataset_item: Dict[str, Any]) -> str:
+        if self.seq_tag_column:
+            seq_tag = dataset_item[self.seq_tag_column]
+            assert isinstance(seq_tag, (str, int, numpy.int64)), f"got {type(seq_tag)} {seq_tag!r}"
+            seq_tag = str(seq_tag)
+        else:
+            seq_tag = f"seq-{corpus_seq_idx}"
+        return seq_tag
+
+    def get_current_seq_order(self) -> Sequence[int]:
+        """:return: list of corpus seq idx"""
+        assert self._seq_order is not None
+        return self._seq_order
+
+    def get_corpus_seq_idx(self, sorted_seq_idx: int) -> int:
+        """:return: corpus seq idx"""
+        return int(self._seq_order[sorted_seq_idx])
+
+
+def get_arrow_shard_files_from_hf_dataset_dir(hf_data_dir: Union[str, os.PathLike]) -> List[str]:
+    """
+    Given some HF datasets directory (created via :func:`datasets.save_to_disk`),
+    return the list of Arrow shard files (``data-*-of-*.arrow``).
+    This also verifies that the directory looks like a valid HF datasets directory.
+    The order of the returned list is by shard index.
+    Note that this does not load the dataset, just inspects the directory structure.
+
+    :param hf_data_dir: directory
+    :return: list of Arrow shard files
+    """
+    hf_data_dir = os.fspath(hf_data_dir)
+    content = os.listdir(hf_data_dir)
+    assert "state.json" in content, f"not a valid HF datasets dir: {hf_data_dir!r}"
+    assert "dataset_info.json" in content, f"not a valid HF datasets dir: {hf_data_dir!r}"
+    pat = re.compile("^(.*)-([0-9]+)-of-([0-9]+).arrow$")
+    content = [pat.match(fn) for fn in content]
+    content = [m for m in content if m]
+    assert content, f"no matching .arrow files in {hf_data_dir!r} found, expected *-*-of-*.arrow"
+    prefix = content[0].group(1)
+    assert all(m.group(1) == prefix for m in content), (
+        f"mismatching prefix in {hf_data_dir!r}, expected {prefix}, got {[m.group(1) for m in content]}"
+    )
+    num_shards = int(content[0].group(3))
+    assert all(int(m.group(3)) == num_shards for m in content), (
+        f"mismatching number of shards in {hf_data_dir!r}, expected {num_shards}, got {[m.group(3) for m in content]}"
+    )
+    assert len(content) == num_shards, f"expected {num_shards} shard files in {hf_data_dir!r}, got {content}"
+    content_by_idx = {int(m.group(2)): m for m in content}
+    assert set(content_by_idx.keys()) == set(range(num_shards)), (
+        f"expected shard indices 0..{num_shards - 1} in {hf_data_dir!r}, got {sorted(content_by_idx.keys())}"
+    )
+    return [hf_data_dir + "/" + content_by_idx[i].group(0) for i in range(num_shards)]
+
+
+def _infer_data_format_for_feature(
+    feature: Union[
+        datasets.features.Sequence,
+        datasets.features.ClassLabel,
+        datasets.features.Value,
+        datasets.features.Array2D,
+        datasets.features.Array3D,
+        datasets.features.Array4D,
+        datasets.features.Audio,
+    ],
+    exc_prefix: str = "",
+) -> Dict[str, Any]:
+    # noinspection PyUnresolvedReferences,PyPackageRequirements
+    import datasets
+
+    labels = None
+    num_classes = None
+    num_dims = 0
+    while isinstance(feature, datasets.features.Sequence):
+        feature: datasets.features.List  # typing for HF datasets 4
+        num_dims += 1
+        if feature.length != -1:
+            num_classes = feature.length
+        feature = feature.feature
+    if isinstance(feature, datasets.features.ClassLabel):
+        labels = feature.names
+        dtype = feature.dtype
+        num_classes = feature.num_classes  # noqa
+    elif isinstance(feature, 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]
+        num_dims += len(feature.shape)
+    elif isinstance(feature, datasets.features.Audio):
+        if feature.decode:
+            dtype = "float32"  # samples
+        else:
+            dtype = "uint8"  # bytes
+        num_dims += 1  # time axis
+    else:
+        assert False, f"{exc_prefix}unsupported feature type {type(feature)} {feature}"
+
+    d = {"dim": num_classes, "ndim": num_dims, "dtype": dtype}
+    if labels:
+        d["sparse"] = True
+        d["vocab"] = {"vocab_file": None, "labels": labels, "unknown_label": None}
+    return d
+
+
+def _make_tensor_template(data: Union[Dict[str, Any], Tensor], name: str) -> Tensor:
+    if isinstance(data, Tensor):
+        data = data.copy(name)
+    else:
+        assert isinstance(data, dict)
+        data = Tensor(name, batch_dim_axis=None, **data)
+    assert data.batch_dim_axis is None
+    return data

{returnn-1.20251006.114241.dist-info → returnn-1.20251007.223754.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251006.114241
+Version: 1.20251007.223754
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20251006.114241.dist-info → returnn-1.20251007.223754.dist-info}/RECORD

@@ -1,9 +1,9 @@
-returnn/PKG-INFO,sha256=khjeqHtAYC68kPEh0ltnHLB14E2_UOObt3MGLvgeMTY,5215
+returnn/PKG-INFO,sha256=lgxotZdSfk01D3LjKSGuuHRSwL3ETFQiSn5GCw80DsE,5215
 returnn/__init__.py,sha256=biBtRsM0WZ406vShaeH-9WFoqJ8XwTbn6g0EeFJ7l8E,1012
 returnn/__main__.py,sha256=lHyZcu_0yc9f7Vf_Kfdy9PmeU0T76XVXnpalHi5WKro,31740
 returnn/__old_mod_loader__.py,sha256=nvsNY-xELdS_IPNkv66Q9Rmvg4dbGW0-EBRDcCmctos,7654
 returnn/__setup__.py,sha256=22kQn2fh11iPM0hLb2Fy5sLmoU1JGvmDxXRYuRgQkwU,4659
-returnn/_setup_info_generated.py,sha256=MFwUI_cce-3opoV_eKfpRd891cRwBo6RMbTGriz-HTg,77
+returnn/_setup_info_generated.py,sha256=kdGINGjXKqjDphnF8IBHiGlKw9_1pozhbIvFUOdv_vU,77
 returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
@@ -13,13 +13,14 @@ returnn/native_op.py,sha256=4_NnvfNxsM8GE_FsD6yOg6PZegqIdtJ3Sl1GdBWmFvg,244424
 returnn/pretrain.py,sha256=MHiXJZqkQFmDVyaYsGpd_Acv20wxl7Pr6s6qJzAT2FI,22648
 returnn/datasets/__init__.py,sha256=PvDlfDOaaopIeUIt0OSvHD2eHZkdkyE-sjMXf35EH5U,390
 returnn/datasets/audio.py,sha256=Gmj7a08dnvYh7Z-G1TNapz42L50AIcDE9JeIZaO1s1M,23334
-returnn/datasets/basic.py,sha256=_42fQztTZq7jNQrWdFBwulB1bNta17LOTyrD8XJ-7_E,73089
+returnn/datasets/basic.py,sha256=s0Vjag5lJ5wGXKENN4KHwGtx7ZDiLdWAFIjFbiqAQsE,74159
 returnn/datasets/bundle_file.py,sha256=KQNrS1MSf-4_idlK0c0KFwON-f5sEK0sWU15WpoMYpE,2380
 returnn/datasets/cached.py,sha256=RyefRjSDdp-HveK-2vLy2C6BIHcpqQ_lNvUKlIa4QAI,25412
 returnn/datasets/cached2.py,sha256=oJOq2lWRQpxm6kyUKW1w5qZBd4kdKEpwM7KY_QnXbq4,11922
-returnn/datasets/distrib_files.py,sha256=-WNVhtvdJFP3L9Meh33oTSYc0FJSvF40mJ5UI_vJbSE,30233
+returnn/datasets/distrib_files.py,sha256=srTieLP02kCepAwZ6Y9p20cqB8nAlVJWbSAoOPna9ik,30567
 returnn/datasets/generating.py,sha256=Qb7V94N_GfL2pZPxWS5PmzszoVXXKzuUmsHuW3dmVbc,99556
 returnn/datasets/hdf.py,sha256=v5sjBenURR9Z-g7AQ9tsL84yDSye5RtbLpym3M6HSDE,67833
+returnn/datasets/huggingface.py,sha256=Bh-1hGYERigvuxjQF8kGwd2gm_BFCPVTtedzk1gz9y0,20042
 returnn/datasets/lm.py,sha256=rQ3jV43lSnlGkKu7m5jTTH7aK0BOMXQocsHfJ8OGec8,99950
 returnn/datasets/map.py,sha256=kOBJVZmwDhLsOplzDNByIfa0NRSUaMo2Lsy36lBvxrM,10907
 returnn/datasets/meta.py,sha256=6XPPxhiNSxWw9Hu5Z6wG8dD9Zk82FqiI-k9HGQSTKgw,95658
@@ -253,8 +254,8 @@ returnn/util/sig_proc.py,sha256=Tjz0VOAVyqu2qDCF5HZ1JjALjcFsHcNkcd96WgZeKfE,7265
 returnn/util/task_system.py,sha256=y4sMVXQ25Qd2z0rx03uOlXlkE-jbCYC1Sjfn-XlraVU,26003
 returnn/util/train_proc_manager.py,sha256=Pjht28k6uz6BNQ47uW6Gf880iyq5q4wx7P_K2tmoAM8,3266
 returnn/util/watch_memory.py,sha256=BR5P2kvBN6UI81cE0_1WAA6Hd1SByLbBaiDxvLhPOew,4213
-returnn-1.20251006.114241.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20251006.114241.dist-info/METADATA,sha256=khjeqHtAYC68kPEh0ltnHLB14E2_UOObt3MGLvgeMTY,5215
-returnn-1.20251006.114241.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20251006.114241.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20251006.114241.dist-info/RECORD,,
+returnn-1.20251007.223754.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20251007.223754.dist-info/METADATA,sha256=lgxotZdSfk01D3LjKSGuuHRSwL3ETFQiSn5GCw80DsE,5215
+returnn-1.20251007.223754.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20251007.223754.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20251007.223754.dist-info/RECORD,,