returnn 1.20250922.155846__tar.gz → 1.20251007.115327__tar.gz

{returnn-1.20250922.155846/returnn.egg-info → returnn-1.20251007.115327}/PKG-INFO

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

returnn-1.20251007.115327/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20251007.115327'
+long_version = '1.20251007.115327+git.70a1d5d'

{returnn-1.20250922.155846 → returnn-1.20251007.115327}/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]]
@@ -1757,3 +1756,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.155846 → returnn-1.20251007.115327}/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-1.20250922.155846 → returnn-1.20251007.115327}/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
@@ -1280,12 +1279,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 +1546,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 +1920,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 +2093,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 +2121,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 +2287,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 +2439,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):
         """

{returnn-1.20250922.155846 → returnn-1.20251007.115327/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250922.155846 → returnn-1.20251007.115327}/tests/test_Dataset.py

@@ -12,7 +12,7 @@ import tempfile
 import contextlib
 from returnn.datasets.generating import Task12AXDataset, DummyDataset, DummyDatasetMultipleSequenceLength
 from returnn.engine.batch import Batch
-from returnn.datasets.basic import Dataset, DatasetSeq, init_dataset
+from returnn.datasets.basic import Dataset, DatasetSeq, init_dataset, _get_seq_len_as_array
 from returnn.util.basic import NumbersDict
 from returnn.util import better_exchook
 
@@ -526,6 +526,26 @@ def test_get_seq_order_laplace_reference():
     assert seq_index == list(seq_index_)
 
 
+def test_get_seq_len_as_array():
+    with tempfile.NamedTemporaryFile(mode="wt", suffix=".npy", encoding="utf8") as tmp_seq_lens_file:
+        num_seqs = 3023
+        rnd = numpy.random.RandomState(42)
+        seq_lens = rnd.randint(1, 23, size=[num_seqs])
+
+        tmp_seq_lens_file.write(repr({f"seq-{i}": int(seq_lens[i]) for i in range(num_seqs)}))
+        tmp_seq_lens_file.write("\n")
+        tmp_seq_lens_file.flush()
+
+        dataset = Dataset()
+        dataset.epoch = 1
+        dataset.seq_ordering = "laplace:.100"
+        dataset._seq_order_seq_lens_file = tmp_seq_lens_file.name
+        dataset.get_all_tags = lambda: [f"seq-{i}" for i in range(num_seqs)]
+
+        seq_lens_ = _get_seq_len_as_array(dataset._get_seq_order_seq_lens_by_idx, num_seqs)
+        assert seq_lens_ is dataset._seq_order_seq_lens_by_idx
+
+
 @contextlib.contextmanager
 def create_ogg_zip_txt_only_dataset_opts(*, text: str = "hello world", seq_tag: str = "sequence0.wav"):
     """create OggZipDataset dict using temp data, consisting of a single sequence with text only"""

returnn-1.20250922.155846/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250922.155846'
-long_version = '1.20250922.155846+git.2334ec6'