returnn 1.20251027.117__py3-none-any.whl → 1.20251027.232712__py3-none-any.whl

returnn/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20251027.117
+Version: 1.20251027.232712
 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.20251027.000117'
-long_version = '1.20251027.000117+git.f3e7971'
+version = '1.20251027.232712'
+long_version = '1.20251027.232712+git.d3f28ed'

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/datasets/generating.py

@@ -1164,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(

returnn/datasets/meta.py

@@ -964,7 +964,6 @@ class CombinedDataset(CachedDataset2):
         self.dataset_keys = set([m[0] for m in data_map.keys()])  # type: typing.Set[str]
         self.dataset_idx2key_map = dict(enumerate(sorted(self.dataset_keys)))  # idx -> dataset-key
         self.data_keys = set(data_map.values())  # type: typing.Set[str]
-        assert "data" in self.data_keys
         self.target_list = sorted(self.data_keys - {"data"})
 
         # Build target lookup table that maps from dataset_key and data_key (data key used by CombinedDataset)
@@ -994,8 +993,7 @@ class CombinedDataset(CachedDataset2):
         if data_dims:
             data_dims = convert_data_dims(data_dims)
             self.data_dims = data_dims
-            assert "data" in data_dims
-            for key in self.target_list:
+            for key in self.data_keys:
                 assert key in data_dims
         else:
             self.data_dims = {}
@@ -1009,7 +1007,7 @@ class CombinedDataset(CachedDataset2):
             if dataset_data_key in dataset.labels:
                 self.labels[data_key] = dataset.labels[dataset_data_key]
 
-        self.num_inputs = self.data_dims["data"][0]
+        self.num_inputs = self.data_dims["data"][0] if "data" in self.data_dims else 0
         self.num_outputs = self.data_dims
 
         self.data_dtypes = {
@@ -1019,6 +1017,9 @@ class CombinedDataset(CachedDataset2):
 
         self.dataset_seq_idx_boundaries: Optional[List[int]] = None
         self.dataset_sorted_seq_idx_list: Optional[List[Tuple[int, int]]] = None
+        self._sub_dataset_cur_loaded_seq_range: Optional[List[Tuple[int, int]]] = None
+        # The usage is about the seqs already covered in dataset_sorted_seq_idx_list,
+        # in case we dynamically build up this list.
         self.used_num_seqs_per_subset: Optional[List[int]] = None
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
@@ -1030,7 +1031,7 @@ class CombinedDataset(CachedDataset2):
         """
 
         assert seq_list is None and seq_order is None, "seq_list and seq_order not supported for %s" % self.__class__
-        need_reinit = self.epoch is None or self.epoch != epoch
+        need_reinit = self.epoch is None or self.epoch != epoch or self.expected_load_seq_start > 0
         num_seqs_saved = self._num_seqs
         super(CombinedDataset, self).init_seq_order(
             epoch=epoch, seq_list=seq_list, seq_order=seq_order
@@ -1047,13 +1048,15 @@ class CombinedDataset(CachedDataset2):
         for dataset in self.datasets.values():
             dataset.init_seq_order(epoch=epoch)
 
+        self._sub_dataset_cur_loaded_seq_range = [(0, 0)] * len(self.datasets)
+
         # noinspection PyBroadException
         try:
             total_num_seqs = sum([self.datasets[k].num_seqs for k in sorted(self.datasets.keys())])
         except Exception:
             total_num_seqs = None
 
-        if total_num_seqs is not None:
+        if total_num_seqs is not None and self.seq_ordering != "interleave":
             self.dataset_seq_idx_boundaries = self._create_dataset_seq_idx_boundaries()
 
             if self.sampling_sizes:
@@ -1090,7 +1093,7 @@ class CombinedDataset(CachedDataset2):
 
             # Re-initialize sequence orders of sub-datasets with created sequence list.
             self.used_num_seqs_per_subset = []
-            for dataset_idx, dataset_key in self.dataset_idx2key_map.items():
+            for dataset_idx, dataset_key in sorted(self.dataset_idx2key_map.items()):
                 assert self.datasets[dataset_key].have_corpus_seq_idx()
                 self.datasets[dataset_key].init_seq_order(epoch=epoch, seq_order=seq_order_subdatasets[dataset_idx])
                 self.used_num_seqs_per_subset.append(len(seq_order_subdatasets[dataset_idx]))
@@ -1098,6 +1101,11 @@ class CombinedDataset(CachedDataset2):
         else:
             self.dataset_sorted_seq_idx_list = []  # We will fill this as we go
             self.used_num_seqs_per_subset = [0] * len(self.datasets)
+            self._num_seqs = total_num_seqs
+
+            # These are currently not supported/implemented.
+            # All of these should just be done in the sub-datasets directly.
+            assert self.partition_epoch == 1 and self.repeat_epoch == 1 and self._num_shards == 1
 
         return True
 
@@ -1236,13 +1244,30 @@ class CombinedDataset(CachedDataset2):
 
         return dataset.get_estimated_seq_length(dataset_seq_idx)
 
-    def _expand_dataset_sec_idxs(self, num_values):
+    def _sub_dataset_make_cur_loaded(self, dataset_idx: int) -> bool:
+        # Cur meaning for the next sequence to be added to dataset_sorted_seq_idx_list.
+        seq_idx = self.used_num_seqs_per_subset[dataset_idx]
+        cur_start, cur_end = self._sub_dataset_cur_loaded_seq_range[dataset_idx]
+        if seq_idx >= cur_end:
+            self._sub_dataset_load_seqs(dataset_idx, cur_start, seq_idx + 1)
+            return True
+        elif seq_idx < cur_start:
+            return False
+        else:
+            return True
+
+    def _expand_dataset_seq_idxs(self, num_values: int) -> bool:
         """
-        :param int num_values: Add num_values entries to the dataset-segment-idx mapping table
-        :return: something?
-        :rtype: bool
+        Try to extend dataset_sorted_seq_idx_list.
+        We expect that we have reached the end of it.
+
+        :param num_values: Add num_values entries to the dataset-segment-idx mapping table
+        :return: whether we added num_values entries
         """
-        for i in range(num_values):
+        for _ in range(num_values):
+            for j in range(len(self.datasets)):
+                self._sub_dataset_make_cur_loaded(j)
+
             if self.seq_ordering == "default":  # i.e. in order
                 dataset_idx = 0
                 while dataset_idx < len(self.datasets):
@@ -1265,6 +1290,32 @@ class CombinedDataset(CachedDataset2):
                 else:
                     return False  # No dataset has remaining data
 
+            elif self.seq_ordering == "interleave":
+                complete_fracs_and_ds_idx = [
+                    (
+                        self.datasets[self.dataset_idx2key_map[j]].get_complete_frac(
+                            self.used_num_seqs_per_subset[j] - 1, allow_only_lr_suitable=True
+                        )
+                        if self.used_num_seqs_per_subset[j] > 0
+                        else 0.0,
+                        j,
+                    )
+                    for j in range(len(self.datasets))
+                ]
+                assert all(frac is not None for frac, _ in complete_fracs_and_ds_idx), (
+                    f"{self}: Datasets must provide complete frac for interleave,"
+                    f" got {complete_fracs_and_ds_idx}, dataset idx2key map {self.dataset_idx2key_map}"
+                )
+                # Sort by complete frac, i.e. datasets with the lowest complete frac first.
+                complete_fracs_and_ds_idx.sort()
+                for complete_frac, dataset_idx in complete_fracs_and_ds_idx:
+                    if self.datasets[self.dataset_idx2key_map[dataset_idx]].is_less_than_num_seqs(
+                        self.used_num_seqs_per_subset[dataset_idx]
+                    ):
+                        break
+                else:
+                    return False  # No dataset has remaining data
+
             elif self.seq_ordering == "random_dataset":
                 while True:
                     # Build probability table
@@ -1323,19 +1374,23 @@ class CombinedDataset(CachedDataset2):
     def _load_seqs(self, start, end):
         # If the segment order is not yet known, fix the next few segments
         if end > len(self.dataset_sorted_seq_idx_list):
-            self._expand_dataset_sec_idxs(end - len(self.dataset_sorted_seq_idx_list))
+            self._expand_dataset_seq_idxs(end - len(self.dataset_sorted_seq_idx_list))
 
         requested_seqs = self.dataset_sorted_seq_idx_list[start:end]
 
         for dataset_idx in range(len(self.datasets)):
-            dataset = self.datasets[self.dataset_idx2key_map[dataset_idx]]
             sub_requested_seqs = [s[1] for s in requested_seqs if s[0] == dataset_idx]
             if not sub_requested_seqs:
                 continue
             sub_start, sub_end = min(sub_requested_seqs), max(sub_requested_seqs)
-            dataset.load_seqs(sub_start, sub_end + 1)
+            self._sub_dataset_load_seqs(dataset_idx, sub_start, sub_end + 1)
         super(CombinedDataset, self)._load_seqs(start=start, end=end)
 
+    def _sub_dataset_load_seqs(self, dataset_idx: int, start: int, end: int):
+        self._sub_dataset_cur_loaded_seq_range[dataset_idx] = (start, end)
+        dataset = self.datasets[self.dataset_idx2key_map[dataset_idx]]
+        dataset.load_seqs(start, end)
+
     def _get_data(self, dataset_key, dataset_seq_idx, data_key):
         """
         :type dataset_seq_idx: int
@@ -1365,19 +1420,30 @@ class CombinedDataset(CachedDataset2):
         dataset = self.datasets[dataset_key]
 
         seq_tag = dataset.get_tag(dataset_seq_idx)
-        features = self._get_data(dataset_key, dataset_seq_idx, "data")
-        targets = {target: self._get_data(dataset_key, dataset_seq_idx, target) for target in self.target_list}
-        return DatasetSeq(seq_idx=seq_idx, seq_tag=seq_tag, features=features, targets=targets)
+        features = {key: self._get_data(dataset_key, dataset_seq_idx, key) for key in self.data_keys}
+        complete_frac = None
+        if self.seq_ordering == "interleave":
+            # In the interleave case, by design, this should be monotonically increasing,
+            # as per how we select the next seq in _expand_dataset_seq_idxs.
+            complete_frac = dataset.get_complete_frac(dataset_seq_idx, allow_only_lr_suitable=True)
+        # In other cases, complete_frac is not so straightforward.
+        # In the case that the total num seqs is known, then it's anyway not necessary.
+        return DatasetSeq(seq_idx=seq_idx, complete_frac=complete_frac, seq_tag=seq_tag, features=features)
 
-    def is_less_than_num_seqs(self, n):
+    def is_less_than_num_seqs(self, n: int) -> bool:
         """
-        :param int n:
-        :rtype: bool
+        :param n:
         """
         if n < len(self.dataset_sorted_seq_idx_list):
             return True
         else:
-            return self._expand_dataset_sec_idxs(n - len(self.dataset_sorted_seq_idx_list) + 1)
+            return self._expand_dataset_seq_idxs(n - len(self.dataset_sorted_seq_idx_list) + 1)
+
+    def get_data_keys(self) -> List[str]:
+        """data keys"""
+        if "data" in self.data_keys:
+            return ["data"] + sorted(self.data_keys - {"data"})
+        return sorted(self.data_keys)
 
     def get_target_list(self):
         """

{returnn-1.20251027.117.dist-info → returnn-1.20251027.232712.dist-info}/METADATA

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

{returnn-1.20251027.117.dist-info → returnn-1.20251027.232712.dist-info}/RECORD

@@ -1,10 +1,10 @@
-returnn/PKG-INFO,sha256=5Pb1iE8plEOp8u6YgK8RC_SSyBmyhFba_D-gcXEE8YI,5212
+returnn/PKG-INFO,sha256=XlAffW31FeRzj4iXwdobRyd-HUqyerhGuIjKnXR-eso,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=KECmOblD-dsBEVI8f_tn-BVnMF4NTy5DuhuYtunMF1M,77
-returnn/config.py,sha256=3tmKhB6FnQZaNdtcYsiB61JnEY--iZ2qmJ4yq0b6tE0,29140
+returnn/_setup_info_generated.py,sha256=mh5Yk4VnansGboCO60Z0keWwnBbHaMw8ywduxfJ0gLM,77
+returnn/config.py,sha256=JK8EjDsUdyY2c90s0KY1rLD1kesVfz6vRT0gxy_AQ5I,29142
 returnn/forward_iface.py,sha256=A_OJiaXsX4MlXQRzST86ylyxSUZbC402PQL1REcqHjM,911
 returnn/learning_rate_control.py,sha256=ZvWryAn_tv9DhV8sh1LV3eE34Yltl3On3mYZAG4hR9s,34684
 returnn/log.py,sha256=WoTDv4XDovgvgXa7iiav-nA8pb25lOEzndbnVrDLfUo,12319
@@ -18,12 +18,12 @@ returnn/datasets/bundle_file.py,sha256=KQNrS1MSf-4_idlK0c0KFwON-f5sEK0sWU15WpoMY
 returnn/datasets/cached.py,sha256=RyefRjSDdp-HveK-2vLy2C6BIHcpqQ_lNvUKlIa4QAI,25412
 returnn/datasets/cached2.py,sha256=oJOq2lWRQpxm6kyUKW1w5qZBd4kdKEpwM7KY_QnXbq4,11922
 returnn/datasets/distrib_files.py,sha256=48edqdf7YpnPJ-TOis3Mz5U9A2DSxfiYT1HCMSti3zw,32718
-returnn/datasets/generating.py,sha256=Qb7V94N_GfL2pZPxWS5PmzszoVXXKzuUmsHuW3dmVbc,99556
+returnn/datasets/generating.py,sha256=o9-JZ2s5QKssux6GcSaM3oivf_PE6nhSOeytRyGB7pQ,99574
 returnn/datasets/hdf.py,sha256=v5sjBenURR9Z-g7AQ9tsL84yDSye5RtbLpym3M6HSDE,67833
 returnn/datasets/huggingface.py,sha256=ls9WMR6gUcMgGksl80g0An1az5Xjya_V3ojbbbsZqrU,20047
 returnn/datasets/lm.py,sha256=rQ3jV43lSnlGkKu7m5jTTH7aK0BOMXQocsHfJ8OGec8,99950
 returnn/datasets/map.py,sha256=kOBJVZmwDhLsOplzDNByIfa0NRSUaMo2Lsy36lBvxrM,10907
-returnn/datasets/meta.py,sha256=E1ZOlIMk4PiNMd5bUCnxdAU7K2hLYEY4Jn6GqbFjjMw,95850
+returnn/datasets/meta.py,sha256=VJ5bk8esq2-b9likNSrCsHQKiLC3Vvti5oBAxg-AsIk,99422
 returnn/datasets/multi_proc.py,sha256=BClXq0fActi1XQa4vcMhHmhYF0Q-fnnDzlIlbBM6_DM,22614
 returnn/datasets/normalization_data.py,sha256=J3njQCMvWAbIAVPepO2L_Xdau9eWYB7Zyd6STeGzTbc,14615
 returnn/datasets/numpy_dump.py,sha256=wl8bKIKAlff2HPJPtuu5wBg3TLOf16d2wLVB4lLAwTM,5158
@@ -255,8 +255,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.20251027.117.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
-returnn-1.20251027.117.dist-info/METADATA,sha256=5Pb1iE8plEOp8u6YgK8RC_SSyBmyhFba_D-gcXEE8YI,5212
-returnn-1.20251027.117.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-returnn-1.20251027.117.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
-returnn-1.20251027.117.dist-info/RECORD,,
+returnn-1.20251027.232712.dist-info/LICENSE,sha256=ywBD_U2aD4vpuoIgNAsjIGBYydl0tVKll3De0Z8s77c,11041
+returnn-1.20251027.232712.dist-info/METADATA,sha256=XlAffW31FeRzj4iXwdobRyd-HUqyerhGuIjKnXR-eso,5215
+returnn-1.20251027.232712.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+returnn-1.20251027.232712.dist-info/top_level.txt,sha256=Lsn4WZc5Pbfk0-xDQOgnFCxOoqxL4CyeM3N1TFbJncw,8
+returnn-1.20251027.232712.dist-info/RECORD,,