returnn 1.20250922.113348__tar.gz → 1.20251104.113756__tar.gz

{returnn-1.20250922.113348/returnn.egg-info → returnn-1.20251104.113756}/PKG-INFO

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

returnn-1.20251104.113756/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20251104.113756'
+long_version = '1.20251104.113756+git.6706c85'

{returnn-1.20250922.113348 → returnn-1.20251104.113756}/returnn/config.py

@@ -801,7 +801,7 @@ class SubProcCopyGlobalConfigPreInitFunc:
         from returnn.log import log
         from returnn import __old_mod_loader__
 
-        better_exchook.install()
+        better_exchook.setup_all()
         __old_mod_loader__.disable_lazy_mod_loads()
 
         if self.global_config:

{returnn-1.20250922.113348 → returnn-1.20251104.113756}/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-1.20250922.113348 → returnn-1.20251104.113756}/returnn/datasets/distrib_files.py

@@ -13,7 +13,7 @@ import sys
 import numpy
 from returnn.log import log
 from returnn.util import better_exchook
-from returnn.util.basic import override_env_var, try_run
+from returnn.util.basic import override_env_var, try_run, OptionalNotImplementedError
 from returnn.util.literal_py_to_pickle import literal_eval
 from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnContext
 from returnn.config import SubProcCopyGlobalConfigPreInitFunc
@@ -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}"
@@ -499,6 +505,24 @@ class DistributeFilesDataset(CachedDataset2):
             self._lazy_init_num_outputs()
         return self._data_keys
 
+    def get_all_tags(self) -> List[str]:
+        """get all tags"""
+        if self.partition_epoch > 1:
+            raise OptionalNotImplementedError(f"{self} get_all_tags not supported for partition_epoch > 1")
+        if self.epoch is None:
+            # Need to init the worker.
+            self.init_seq_order(epoch=1)
+        return self._workers[self.epoch].get_all_tags()
+
+    def get_total_num_seqs(self, *, fast: bool = False) -> int:
+        """get total num seqs"""
+        if self.partition_epoch > 1:
+            raise OptionalNotImplementedError(f"{self} get_total_num_seqs not supported for partition_epoch > 1")
+        if self.epoch is None:
+            # Need to init the worker.
+            self.init_seq_order(epoch=1)
+        return self._workers[self.epoch].get_total_num_seqs(fast=fast)
+
 
 def _get_key_for_file_tree(t: FileTree) -> str:
     """generates a deterministic key given a file tree"""
@@ -602,6 +626,26 @@ class _WorkerProcParent:
         assert msg == "data_seq"
         return data
 
+    def get_all_tags(self) -> List[str]:
+        """get all tags"""
+        self._lazy_wait_for_init_seq_order()
+        self.parent_conn.send(("get_all_tags", {}))
+        msg, data = self.parent_conn.recv()
+        assert msg == "all_tags"
+        if isinstance(data, Exception):
+            raise data
+        return data
+
+    def get_total_num_seqs(self, **kwargs) -> int:
+        """get total num seqs"""
+        self._lazy_wait_for_init_seq_order()
+        self.parent_conn.send(("get_total_num_seqs", kwargs))
+        msg, data = self.parent_conn.recv()
+        assert msg == "total_num_seqs"
+        if isinstance(data, Exception):
+            raise data
+        return data
+
     def exit(self, *, join: bool = True):
         """exit"""
         self._lazy_wait_for_init_seq_order()
@@ -716,6 +760,20 @@ def _worker_proc_loop(
                 got_init_seq_order = True
                 next_seq_idx = 0
                 cache.clear()
+            elif msg == "get_all_tags":
+                try:
+                    tags = dataset.get_all_tags()
+                except Exception as exc:
+                    parent_conn.send(("all_tags", exc))
+                else:
+                    parent_conn.send(("all_tags", tags))
+            elif msg == "get_total_num_seqs":
+                try:
+                    total_num_seqs = dataset.get_total_num_seqs(**kwargs)
+                except Exception as exc:
+                    parent_conn.send(("total_num_seqs", exc))
+                else:
+                    parent_conn.send(("total_num_seqs", total_num_seqs))
             else:
                 raise Exception(f"unknown msg {msg!r}")
     except KeyboardInterrupt:  # when parent dies

{returnn-1.20250922.113348 → returnn-1.20251104.113756}/returnn/datasets/generating.py

@@ -7,14 +7,13 @@ from __future__ import annotations
 from typing import Optional, Union, Any, Sequence, List, Tuple, Dict
 import numpy
 import sys
-import typing
 
 from returnn.util.basic import class_idx_seq_to_1_of_k, CollectionReadCheckCovered
 from returnn.log import log
 from returnn.tensor import Tensor, Dim, TensorDict
 
 from .util.feature_extraction import ExtractAudioFeatures
-from .util.vocabulary import *
+from .util.vocabulary import Vocabulary, BytePairEncoding, CharacterTargets
 from .audio import OggZipDataset  # noqa # for API compatibility
 from .basic import Dataset, DatasetSeq, convert_data_dims
 from .cached2 import CachedDataset2
@@ -1165,11 +1164,9 @@ class StaticDataset(CachedDataset2):
         """supports sorting"""
         return True
 
-    def _collect_single_seq(self, seq_idx):
-        """
-        :param int seq_idx:
-        :rtype: DatasetSeq
-        """
+    def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
+        if seq_idx >= len(self._seq_order):
+            return None
         corpus_seq_idx = self._seq_order[seq_idx]
         data = self.data[corpus_seq_idx]
         return DatasetSeq(
@@ -1280,12 +1277,6 @@ class CopyTaskDataset(GeneratingDataset):
         return DatasetSeq(seq_idx=seq_idx, features=seq_np, targets={"classes": seq_np})
 
 
-# Multiple external sources where we could write automatic wrappers:
-# * https://github.com/tensorflow/datasets
-# * tf.contrib.keras.datasets, https://www.tensorflow.org/api_docs/python/tf/keras/datasets
-# * nltk.corpus
-
-
 class TimitDataset(CachedDataset2):
     """
     DARPA TIMIT Acoustic-Phonetic Continuous Speech Corpus.
@@ -1553,7 +1544,7 @@ class TimitDataset(CachedDataset2):
 
         self._random_permute_audio = CollectionReadCheckCovered.from_bool_or_dict(random_permute_audio)
 
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
         self._init_timit()
 
         self._audio_data = {}  # seq_tag -> (audio, sample_rate). loaded by self._reader_thread_main
@@ -1927,8 +1918,8 @@ class BlissDataset(CachedDataset2):
         self._with_delta = with_delta
         self.num_inputs *= 1 + with_delta
         self._bpe_file = open(bpe_file, "r")
-        self._seqs = []  # type: typing.List[BlissDataset.SeqInfo]
-        self._vocab = {}  # type: typing.Dict[str,int]  # set in self._parse_vocab
+        self._seqs: List[BlissDataset.SeqInfo] = []
+        self._vocab: Dict[str, int] = {}  # set in self._parse_vocab
         self._parse_bliss_xml(filename=path)
         # TODO: loading audio like in TimitDataset, and in parallel
         self._bpe = BytePairEncoding(vocab_file=vocab_file, bpe_file=bpe_file)
@@ -2100,7 +2091,7 @@ class LibriSpeechCorpus(CachedDataset2):
             self.targets = CharacterTargets(**chars)
         elif targets is None:
             assert bpe is None and chars is None
-            self.targets = None  # type: typing.Optional[Vocabulary]
+            self.targets: Optional[Vocabulary] = None
         else:
             raise Exception("invalid targets %r. provide bpe or chars" % targets)
         if self.targets:
@@ -2128,7 +2119,7 @@ class LibriSpeechCorpus(CachedDataset2):
             self._reference_seq_order = seqs
             self.transs = {s: self.transs[s] for s in seqs}
         self.epoch_wise_filter = epoch_wise_filter
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
         self.init_seq_order()
 
     def _collect_trans(self):
@@ -2294,9 +2285,9 @@ class LibriSpeechCorpus(CachedDataset2):
         """:return: whether this dataset supports sharding"""
         return True
 
-    def get_current_seq_order(self):
+    def get_current_seq_order(self) -> Sequence[int]:
         """
-        :rtype: typing.Sequence[int]
+        :return: seq order of current epoch
         """
         assert self._seq_order is not None
         return self._seq_order
@@ -2446,7 +2437,7 @@ class Enwik8Corpus(CachedDataset2):
         self._batch_num_seqs = batch_num_seqs
         self._random = numpy.random.RandomState(1)  # seed will be set in init_seq_order
         self._seq_starts = numpy.arange(0, len(self._data) - 1, seq_len)
-        self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
+        self._seq_order: Optional[Sequence[int]] = None
 
     def get_data_dtype(self, key):
         """