returnn 1.20250227.110407__tar.gz → 1.20251105.112921__tar.gz

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/CONTRIBUTING.md

@@ -18,7 +18,7 @@ General rules when contributing to the code of RETURNN:
   Our code style uses most common Python conventions.
   If you are not an expert in Python, use PyCharm,
   and follow [our PyCharm configuration guide](https://github.com/rwth-i6/returnn/wiki/PyCharm-Configuration).
-  Apply [black](https://black.readthedocs.io/).
+  Apply [ruff](https://github.com/astral-sh/ruff).
 * Make sure all [tests](https://returnn.readthedocs.io/en/latest/advanced/test_suite.html) pass.
 * At the time being, we want to support earlier versions of TF 1
   (consider at least TF 1.8, but maybe even TF 1.4)

{returnn-1.20250227.110407/returnn.egg-info → returnn-1.20251105.112921}/PKG-INFO

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

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/__init__.py

@@ -7,7 +7,6 @@ This was used for the old flat code file structure, for usage like::
 We want to support the same code.
 """
 
-
 from __future__ import annotations
 import os
 import sys

returnn-1.20251105.112921/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20251105.112921'
+long_version = '1.20251105.112921+git.d5d3efd'

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/demos/demo-tf-search-compiled-graph.py

@@ -8,12 +8,12 @@ This is just for demonstration, testing and debugging purpose. The search itself
 
 # No RETURNN dependency needed for the basic search. Just TF itself.
 
-import typing
 import os
 import json
 import argparse
 import tensorflow as tf
 import numpy
+from typing import List, Optional, Tuple
 
 
 class Hyp:
@@ -26,7 +26,7 @@ class Hyp:
         :param int idx: hyp idx (to identify it in a beam)
         """
         self.idx = idx
-        self.source_idx = None  # type: typing.Optional[int]  # source hyp idx
+        self.source_idx: Optional[int] = None  # source hyp idx
         self.score = 0.0
         self.seq = []  # label seq
 
@@ -91,7 +91,6 @@ def main():
         # Now loop over decoder steps.
         max_dec_len = 100  # TODO better default... depending on input len. or configurable...
         for i in range(max_dec_len):
-
             # Loop over all stochastic variables.
             for stochastic_var in info["stochastic_var_order"]:
                 assert isinstance(stochastic_var, str)
@@ -108,9 +107,7 @@ def main():
                 # TODO: length norm here?
 
                 # Select new hypotheses.
-                best_possibilities = sorted(all_possibilities)[
-                    : args.beam_size
-                ]  # type: typing.List[typing.Tuple[float,int,Hyp]]
+                best_possibilities: List[Tuple[float, int, Hyp]] = sorted(all_possibilities)[: args.beam_size]
                 assert len(best_possibilities) == args.beam_size
                 hyps = [
                     hyp.expand(idx=i, label=label, score=score)
@@ -121,8 +118,9 @@ def main():
                 session.run(
                     info["state_vars"]["stochastic_var_scores_%s" % stochastic_var] + "/Assign...?",  # TODO...
                     feed_dict={
-                        info["state_vars"]["stochastic_var_scores_%s" % stochastic_var]
-                        + "/Initial...?": [[hyp.seq[-1] for hyp in hyps]]  # TODO...
+                        info["state_vars"]["stochastic_var_scores_%s" % stochastic_var] + "/Initial...?": [
+                            [hyp.seq[-1] for hyp in hyps]
+                        ]  # TODO...
                     },
                 )
 

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/pyproject.toml

@@ -13,5 +13,9 @@ extend-exclude = '''
 )/
 '''
 
+[tool.ruff]
+line-length = 120
+target-version = "py38"  # https://github.com/rwth-i6/returnn/issues/1326
+
 [build-system]
 requires = ["setuptools", "numpy"]

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/returnn/__main__.py

@@ -34,21 +34,21 @@ from returnn.util.basic import BackendEngine, BehaviorVersion
 
 # These imports are not directly used here, but make them available, as other code imports them from here.
 # noinspection PyUnresolvedReferences
-from returnn.util.debug import init_ipython_kernel, init_better_exchook, init_faulthandler, debug_shell
+from returnn.util.debug import init_ipython_kernel, init_better_exchook, init_faulthandler, debug_shell  # noqa: F401
 
 # Some external scripts import those functions from here, thus keep this here.
 # noinspection PyUnresolvedReferences
-from returnn.util.basic import init_thread_join_hack, describe_returnn_version
+from returnn.util.basic import init_thread_join_hack, describe_returnn_version  # noqa: F401
 
 if TYPE_CHECKING:
     import returnn.tf.engine
     import returnn.torch.engine
 
-config = None  # type: Optional[Config]
-engine = None  # type: Optional[Union[returnn.tf.engine.Engine, returnn.torch.engine.Engine]]
-train_data = None  # type: Optional[Dataset]
-dev_data = None  # type: Optional[Dataset]
-eval_data = None  # type: Optional[Dataset]
+config: Optional[Config] = None
+engine: Optional[Union[returnn.tf.engine.Engine, returnn.torch.engine.Engine]] = None
+train_data: Optional[Dataset] = None
+dev_data: Optional[Dataset] = None
+eval_data: Optional[Dataset] = None
 quit_returnn = False
 
 

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/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.20250227.110407 → returnn-1.20251105.112921}/returnn/datasets/basic.py

@@ -19,8 +19,8 @@ import os
 import math
 import numpy
 import functools
-import typing
-from typing import TYPE_CHECKING, Optional, Any, Union, Type, Dict, Sequence, List, Callable
+import types
+from typing import TYPE_CHECKING, Optional, Any, Set, Tuple, Union, Type, Dict, Sequence, List, Callable
 
 from returnn.log import log
 from returnn.engine.batch import Batch, BatchSetGenerator
@@ -141,12 +141,10 @@ class Dataset:
         :param int _shard_index: local shard index, when sharding is enabled
         """
         self.name = name or ("dataset_id%s" % id(self))
-        self.lock = None  # type: Optional[RLock]  # Used when manipulating our data potentially from multiple threads.
-        self.rnd_seq_drop = None  # type: typing.Optional[Random]
+        self.lock: Optional[RLock] = None  # Used when manipulating our data potentially from multiple threads.
+        self.rnd_seq_drop: Optional[Random] = None
         self.num_inputs = 0  # usually not used, but num_outputs instead, which is more generic
-        self.num_outputs = (
-            None
-        )  # type: typing.Optional[typing.Dict[str,typing.Tuple[int,int]]]  # tuple is num-classes, len(shape).  # nopep8
+        self.num_outputs: Optional[Dict[str, Tuple[int, int]]] = None  # tuple is num-classes, len(shape).
         self.window = window
         self.seq_ordering = seq_ordering  # "default", "sorted" or "random". See self.get_seq_order_for_epoch().
         self.fixed_random_seed = fixed_random_seed
@@ -157,12 +155,12 @@ 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."
-        self.labels = {}  # type: typing.Dict[str,typing.List[str]]
+        assert self.partition_epoch == 1 or self.repeat_epoch == 1, (
+            "Combining partition_epoch and repeat_epoch is prohibited."
+        )
+        self.labels: Dict[str, List[str]] = {}
         self.weights = {}
         self._num_timesteps = 0
         self._num_seqs = 0
@@ -213,8 +211,8 @@ class Dataset:
             getattr(self, "epoch", "<unknown>"),
         )
 
-    _getnewargs_exclude_attrs = set()  # type: typing.Set[str]
-    _getnewargs_remap = {}  # type: typing.Dict[str,str]
+    _getnewargs_exclude_attrs: Set[str] = set()
+    _getnewargs_remap: Dict[str, str] = {}
 
     @staticmethod
     def _create_from_reduce(cls, kwargs, state) -> Dataset:
@@ -430,9 +428,9 @@ class Dataset:
         """
         raise OptionalNotImplementedError
 
-    def get_num_timesteps(self):
+    def get_num_timesteps(self) -> Union[int, NumbersDict]:
         """
-        :rtype: int
+        :return: how much frames we have in total.
         """
         assert self._num_timesteps > 0
         return self._num_timesteps
@@ -489,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
@@ -505,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.
@@ -518,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,20 +556,33 @@ class Dataset:
             for i in range(1, num):
                 seq_index[i::num] += i * (num_seqs // num)
         elif seq_ordering_method == "reverse":
-            seq_index = range(num_seqs - 1, -1, -1)  # type: Union[range, typing.Sequence[int]]
+            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.startswith("random"):
-            tmp = seq_ordering_method.split(":")
+        elif seq_ordering_method == "random" or seq_ordering_method.startswith("random:"):
+            tmp = seq_ordering_method.split(":", 1)
             nth = int(tmp[1]) if len(tmp) > 1 else 1
             # Keep this deterministic! Use fixed seed.
             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)
-        elif seq_ordering_method.startswith("sort_bin_shuffle"):
+        elif seq_ordering_method == "random_sample" or seq_ordering_method.startswith("random_sample:"):
+            tmp = seq_ordering_method.split(":", 1)
+            nth = int(tmp[1]) if len(tmp) > 1 else 1
+            # Keep this deterministic! Use fixed seed.
+            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.randint(0, num_seqs, size=num_seqs)
+        elif (
+            seq_ordering_method == "sort_bin_shuffle"
+            or seq_ordering_method.startswith("sort_bin_shuffle:")
+            or seq_ordering_method == "sort_bin_shuffle_x2"
+            or seq_ordering_method.startswith("sort_bin_shuffle_x2:")
+        ):
             # Shuffle seqs, sort by length, and shuffle bins (then shuffle seqs within each bin if sort_bin_shuffle_x2).
             assert get_seq_len
             tmp = seq_ordering_method.split(":")[1:]
@@ -605,7 +613,7 @@ class Dataset:
                     random_generator.shuffle(part)  # Shuffle within the bin.
                 out_index.append(part)
             seq_index = numpy.concatenate(out_index)
-        elif seq_ordering_method.startswith("laplace"):
+        elif seq_ordering_method == "laplace" or seq_ordering_method.startswith("laplace:"):
             assert get_seq_len
             tmp = seq_ordering_method.split(":")[1:]
             if len(tmp) == 0:
@@ -619,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]]
@@ -660,12 +668,13 @@ class Dataset:
             )
             old_seq_index = seq_index
             seq_index = [i for i in seq_index if all_seq_tags[i] in self.seq_tags_filter]
-            assert (
-                seq_index
-            ), "%s: empty after applying seq_list_filter_file. Example filter tags: %r, used tags: %r" % (
-                self,
-                sorted(self.seq_tags_filter)[:3],
-                [all_seq_tags[i] for i in old_seq_index[:3]],
+            assert seq_index, (
+                "%s: empty after applying seq_list_filter_file. Example filter tags: %r, used tags: %r"
+                % (
+                    self,
+                    sorted(self.seq_tags_filter)[:3],
+                    [all_seq_tags[i] for i in old_seq_index[:3]],
+                )
             )
         return seq_index
 
@@ -736,9 +745,9 @@ class Dataset:
         """
         self.epoch = epoch
         self.rnd_seq_drop = Random(self._get_random_seed_for_epoch(epoch=epoch))
-        assert (
-            self._num_shards == 1 or self.supports_sharding()
-        ), f"{self}: does not support sharding, but got num_shards == {self._num_shards}"
+        assert self._num_shards == 1 or self.supports_sharding(), (
+            f"{self}: does not support sharding, but got num_shards == {self._num_shards}"
+        )
         return False
 
     def finish_epoch(self, *, free_resources: bool = False):
@@ -749,12 +758,11 @@ class Dataset:
         """
         self.epoch = None
 
-    def get_current_seq_order(self):
+    def get_current_seq_order(self) -> Sequence[int]:
         """
         :return: many datasets use self.get_seq_order_for_epoch. this function would return the current seq order
           for the current epoch, after self.init_seq_order was called.
           Not all datasets implement this.
-        :rtype: typing.Sequence[int]
         """
         raise OptionalNotImplementedError
 
@@ -864,18 +872,16 @@ class Dataset:
             data = self.get_data(seq_idx, key)
             return data[s0_start:s0_end]
 
-    def get_tag(self, sorted_seq_idx):
+    def get_tag(self, sorted_seq_idx: int) -> str:
         """
-        :param int sorted_seq_idx:
-        :rtype: str
+        :param sorted_seq_idx:
         """
         return "seq-%i" % sorted_seq_idx
 
-    def get_all_tags(self):
+    def get_all_tags(self) -> List[str]:
         """
         :return: list of all seq tags, of the whole dataset, without partition epoch.
           Note that this is not possible with all datasets.
-        :rtype: list[str]
         """
         raise OptionalNotImplementedError(f"{self} get_all_tags not implemented")
 
@@ -905,7 +911,7 @@ class Dataset:
         if self.seq_ordering == "default" and self.partition_epoch == 1:
             return seq_idx
         assert self.have_corpus_seq_idx()
-        raise NotImplemented
+        raise NotImplementedError
 
     def have_get_corpus_seq(self) -> bool:
         """
@@ -972,16 +978,16 @@ class Dataset:
             except Exception:  # also not always available
                 num_seqs = None  # ignore
 
-        if math.isinf(num_seqs):
+        if num_seqs is not None and math.isinf(num_seqs):
             if allow_only_lr_suitable:
                 # cannot compute meaningful complete_frac for infinite num_seqs
                 return None
             else:
                 num_seqs = None
 
-        assert (
-            num_seqs is None or 0 <= sorted_seq_idx < num_seqs
-        ), f"{self}: invalid seq indices: 0 <= seq_idx ({sorted_seq_idx}) < num_seqs ({num_seqs}) violated"
+        assert num_seqs is None or 0 <= sorted_seq_idx < num_seqs, (
+            f"{self}: invalid seq indices: 0 <= seq_idx ({sorted_seq_idx}) < num_seqs ({num_seqs}) violated"
+        )
         return self.generic_complete_frac(sorted_seq_idx, num_seqs)
 
     @property
@@ -1064,7 +1070,7 @@ class Dataset:
         if key in self.num_outputs:
             if self.num_outputs[key][1] <= 1:
                 return []
-            res_shape = [None] * (self.num_outputs[key][1] - 1)  # type: typing.List[typing.Union[None,int]]
+            res_shape: List[Union[None, int]] = [None] * (self.num_outputs[key][1] - 1)
             if not self.is_data_sparse(key):
                 res_shape[-1] = self.get_data_dim(key)
             return res_shape
@@ -1118,7 +1124,9 @@ class Dataset:
 
     def serialize_data(self, key: str, data: numpy.ndarray) -> str:
         """
-        In case you have a :class:`Vocabulary`, just use :func:`Vocabulary.get_seq_labels`.
+        This is deprecated, as this is slow!
+        In case you have a :class:`Vocabulary`, just use :func:`Vocabulary.get_seq_labels`
+        or :func:`Vocabulary.serialize_labels`.
 
         :param key: e.g. "classes". self.labels[key] should be set
         :param numpy.ndarray data: 0D or 1D
@@ -1492,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)
@@ -1748,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.20250227.110407 → returnn-1.20251105.112921}/returnn/datasets/cached.py

@@ -46,9 +46,10 @@ class CachedDataset(Dataset):
         self._index_map = range(len(self._seq_index))  # sorted seq idx -> seq_index idx
         self._tag_idx = {}  # type: typing.Dict[str,int]  # map of tag -> real-seq-idx. call _update_tag_idx
         self.targets = {}
-        self.target_keys = (
-            []
-        )  # the keys for which we provide data; we may have labels for additional keys in self.labels
+        # the keys for which we provide data;
+        # we may have labels for additional keys in self.labels
+        self.target_keys = []
+
         self.timestamps = None
 
     def initialize(self):

{returnn-1.20250227.110407 → returnn-1.20251105.112921}/returnn/datasets/cached2.py

@@ -4,18 +4,10 @@ Provides :class:`CachedDataset2`.
 
 from __future__ import annotations
 import numpy
-import typing
-from typing import Optional
+from typing import Optional, List
 from threading import Condition
 from .basic import Dataset, DatasetSeq
 
-try:
-    # noinspection PyCompatibility
-    from _thread import interrupt_main
-except ImportError:
-    # noinspection PyUnresolvedReferences,PyCompatibility
-    from thread import interrupt_main
-
 
 class CachedDataset2(Dataset):
     """
@@ -36,7 +28,7 @@ class CachedDataset2(Dataset):
         self._num_timesteps = None
         self.epoch = None
         self.reached_final_seq = False
-        self.added_data = []  # type: typing.List[DatasetSeq]
+        self.added_data: List[DatasetSeq] = []
         self.expected_load_seq_start = 0
         self._num_timesteps_accumulated = 0