returnn 1.20250901.123052__py3-none-any.whl → 1.20260105.192646__py3-none-any.whl

returnn/datasets/meta.py

@@ -253,22 +253,12 @@ class MetaDataset(CachedDataset2):
         }
 
         self._seq_list_file = seq_list_file
-        self.seq_list_original = self._load_seq_list(seq_list_file)
-        self.num_total_seqs = len(self.seq_list_original[self.default_dataset_key])
-        for key in self.dataset_keys:
-            assert len(self.seq_list_original[key]) == self.num_total_seqs
-
-        self.tag_idx = {tag: idx for (idx, tag) in enumerate(self.seq_list_original[self.default_dataset_key])}
+        self.seq_list_original: Optional[Dict[str, List[str]]] = None
+        self.tag_idx: Optional[Dict[str, int]] = None
 
         self._seq_lens: Optional[Dict[str, NumbersDict]] = None
         self._num_timesteps: Optional[NumbersDict] = None
         self._seq_lens_file = seq_lens_file
-        if seq_lens_file:
-            seq_lens = load_json(filename=seq_lens_file)
-            assert isinstance(seq_lens, dict)
-            # dict[str,NumbersDict], seq-tag -> data-key -> len
-            self._seq_lens = {tag: NumbersDict(l) for (tag, l) in seq_lens.items()}
-            self._num_timesteps = sum([self._seq_lens[s] for s in self.seq_list_original[self.default_dataset_key]])
 
         if data_dims:
             data_dims = convert_data_dims(data_dims)
@@ -290,19 +280,20 @@ class MetaDataset(CachedDataset2):
         self.num_outputs = self.data_dims
 
         self.orig_seq_order_is_initialized = False
+        self._current_seq_order: List[int] = []
         self.seq_list_ordered: Optional[Dict[str, List[str]]] = None
 
-    def _load_seq_list(self, seq_list_file: Optional[Union[str, Dict[str, str]]] = None) -> Dict[str, List[str]]:
-        """
-        :param seq_list_file:
-        :return: dict: dataset key -> seq list
-        """
-        if not seq_list_file:
+    def _lazy_init_seq_list(self):
+        if self.seq_list_original is not None:
+            return
+
+        if not self._seq_list_file:
             # We create a sequence list from all the sequences of the default dataset
             # and hope that it also applies to the
             # other datasets.
             # This can only work if all datasets have the same tag format and the sequences in the other
             # datasets are a subset of those in the default dataset.
+            # (But the order does not matter.)
             default_dataset = self.datasets[self.default_dataset_key]
             assert isinstance(default_dataset, Dataset)
             print(
@@ -349,17 +340,18 @@ class MetaDataset(CachedDataset2):
                         break  # only print one
                 del seq_list_set
                 raise Exception("Dataset %r is missing seqs." % key)
-        elif isinstance(seq_list_file, str):
-            seq_list = Dataset._load_seq_list_file(seq_list_file, expect_list=False)
-        elif isinstance(seq_list_file, dict):
+        elif isinstance(self._seq_list_file, str):
+            seq_list = Dataset._load_seq_list_file(self._seq_list_file, expect_list=False)
+        elif isinstance(self._seq_list_file, dict):
             for key in self.dataset_keys:
-                if key not in seq_list_file:
+                if key not in self._seq_list_file:
                     raise ValueError(f"seq_list_file does not contain all datasets, missing {key}")
-            seq_list = {key: Dataset._load_seq_list_file(seq_list_file[key]) for key in self.dataset_keys}
+            seq_list = {key: Dataset._load_seq_list_file(self._seq_list_file[key]) for key in self.dataset_keys}
         else:
-            raise TypeError(f"unexpected seq_list_file type {type(seq_list_file)}")
+            raise TypeError(f"unexpected seq_list_file type {type(self._seq_list_file)}")
 
         if isinstance(seq_list, list):
+            # Use same seq list for all datasets
             seq_list = {key: seq_list for key in self.dataset_keys}
         elif isinstance(seq_list, dict):
             for key in self.dataset_keys:
@@ -368,10 +360,29 @@ class MetaDataset(CachedDataset2):
         else:
             raise TypeError(f"unexpected seq_list type {type(seq_list)}")
 
-        return seq_list
+        for key in self.dataset_keys:
+            assert len(seq_list[key]) == len(seq_list[self.default_dataset_key])
+
+        self.seq_list_original = seq_list
+
+    def _lazy_init_tag_idx(self):
+        if self.tag_idx is not None:
+            return
+        self._lazy_init_seq_list()
+        self.tag_idx = {tag: idx for (idx, tag) in enumerate(self.seq_list_original[self.default_dataset_key])}
+
+    def _lazy_init_seq_lens(self):
+        if self._seq_lens is not None:
+            return
+        assert self._seq_lens_file
+        seq_lens = load_json(filename=self._seq_lens_file)
+        assert isinstance(seq_lens, dict)
+        # dict[str,NumbersDict], seq-tag -> data-key -> len
+        self._seq_lens = {tag: NumbersDict(lens) for (tag, lens) in seq_lens.items()}
 
     def _get_dataset_seq_length(self, seq_idx: int):
         if not self.orig_seq_order_is_initialized:
+            self._lazy_init_seq_list()
             # To use get_seq_length() we first have to init the sequence order once in original order.
             # If sequence lengths are not needed by get_seq_order_for_epoch this is never executed.
             self.datasets[self.default_dataset_key].init_seq_order(
@@ -379,6 +390,9 @@ class MetaDataset(CachedDataset2):
             )
             self.orig_seq_order_is_initialized = True
 
+        # Warning: This is not correct in the general case.
+        # get_seq_length needs to have load_seqs called beforehand per API contract.
+        # For some datasets, it might anyway work.
         return self.datasets[self.default_dataset_key].get_seq_length(seq_idx)["data"]
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
@@ -392,6 +406,7 @@ class MetaDataset(CachedDataset2):
             self.epoch is None
             or self.epoch != epoch
             or self.seq_list_ordered is None
+            or not self._current_seq_order
             or seq_list is not None
             or seq_order is not None
             or self.expected_load_seq_start > 0
@@ -401,16 +416,17 @@ class MetaDataset(CachedDataset2):
             # This is called via initialize() with epoch=None, just to init some other things.
             # We are not expected to have prepared any real epoch here.
             self._num_seqs = 0
+            self._current_seq_order = []
             return True
 
         if not need_reinit:
-            self._num_seqs = len(self.seq_list_ordered[self.default_dataset_key])
             return False
 
         seq_order_dataset = None
         if seq_order is not None:
             seq_index = seq_order
         elif seq_list is not None:
+            self._lazy_init_tag_idx()
             seq_index = [self.tag_idx[tag] for tag in seq_list]
         elif self.seq_order_control_dataset:
             seq_order_dataset = self.datasets[self.seq_order_control_dataset]
@@ -418,13 +434,15 @@ class MetaDataset(CachedDataset2):
             seq_order_dataset.init_seq_order(epoch=epoch)
             seq_index = seq_order_dataset.get_current_seq_order()
         else:
-            if self._seq_lens:
+            if self._seq_lens_file:
 
                 def get_seq_len(s):
                     """
                     :param int s:
                     :rtype: int
                     """
+                    self._lazy_init_seq_list()
+                    self._lazy_init_seq_lens()
                     return self._seq_lens[self.seq_list_original[self.default_dataset_key][s]]["data"]
 
             elif self._seq_order_seq_lens_file:
@@ -432,8 +450,10 @@ class MetaDataset(CachedDataset2):
             else:
                 self.orig_seq_order_is_initialized = False
                 get_seq_len = self._get_dataset_seq_length
-            seq_index = self.get_seq_order_for_epoch(epoch, self.num_total_seqs, get_seq_len)
+            seq_index = self.get_seq_order_for_epoch(epoch, self.get_total_num_seqs(), get_seq_len)
         self._num_seqs = len(seq_index)
+        self._current_seq_order = seq_index
+        self._lazy_init_seq_list()
         self.seq_list_ordered = {key: [ls[s] for s in seq_index] for (key, ls) in self.seq_list_original.items()}
 
         for dataset_key, dataset in self.datasets.items():
@@ -447,7 +467,7 @@ class MetaDataset(CachedDataset2):
         """supports sorting"""
         if self.seq_order_control_dataset:
             return self.datasets[self.seq_order_control_dataset].supports_seq_order_sorting()
-        if self._seq_lens or self._seq_order_seq_lens_file:
+        if self._seq_lens_file or self._seq_order_seq_lens_file:
             return True
         return False
 
@@ -464,20 +484,40 @@ class MetaDataset(CachedDataset2):
         :return: current seq order for the current epoch, after self.init_seq_order was called.
         :rtype: list[int]
         """
-        return [self.tag_idx[tag] for tag in self.seq_list_ordered[self.default_dataset_key]]
+        return self._current_seq_order
 
     def get_all_tags(self):
         """
         :return: list of all seq tags, of the whole dataset, without partition epoch
         :rtype: list[str]
         """
+        if self._seq_list_file is None:
+            return self.datasets[self.default_dataset_key].get_all_tags()
+        self._lazy_init_seq_list()
+        assert self.seq_list_original is not None
         return self.seq_list_original[self.default_dataset_key]
 
     def get_total_num_seqs(self, *, fast: bool = False) -> int:
         """
+        :param fast: if True, might raise an exception if not possible to get fast.
         :return: total number of seqs, without partition epoch
         """
-        return self.num_total_seqs
+        if self._seq_list_file is None:
+            return self.datasets[self.default_dataset_key].get_total_num_seqs(fast=fast)
+        if fast and self.seq_list_original is None:
+            raise OptionalNotImplementedError(f"{self} get_total_num_seqs, seq list not loaded yet")
+        self._lazy_init_seq_list()
+        assert self.seq_list_original is not None
+        return len(self.seq_list_original[self.default_dataset_key])
+
+    def get_num_timesteps(self):
+        """num timesteps"""
+        if self._num_timesteps is None and self._seq_lens_file:
+            self._lazy_init_seq_lens()
+            self._num_timesteps = sum([self._seq_lens[s] for s in self.get_all_tags()], start=NumbersDict())
+        if self._seq_list_file is None:
+            return self.datasets[self.default_dataset_key].get_num_timesteps()
+        return super().get_num_timesteps()
 
     def finish_epoch(self, *, free_resources: bool = False):
         """
@@ -503,8 +543,9 @@ class MetaDataset(CachedDataset2):
         if start_ < end:
             for dataset_key in self.dataset_keys:
                 self.datasets[dataset_key].load_seqs(start_, end)
-                for seq_idx in range(start_, end):
-                    self._check_dataset_seq(dataset_key, seq_idx)
+                if self.seq_list_ordered is not None:
+                    for seq_idx in range(start_, end):
+                        self._check_dataset_seq(dataset_key, seq_idx)
         super(MetaDataset, self)._load_seqs(start=start, end=end)
 
     def _check_dataset_seq(self, dataset_key, seq_idx):
@@ -531,7 +572,7 @@ class MetaDataset(CachedDataset2):
         :type seq_idx: int
         :rtype: DatasetSeq
         """
-        seq_tag = self.seq_list_ordered[self.default_dataset_key][seq_idx]
+        seq_tag = self.get_tag(seq_idx)
         features = {data_key: self._get_data(seq_idx, data_key) for data_key in self.data_keys}
         return DatasetSeq(seq_idx=seq_idx, seq_tag=seq_tag, features=features)
 
@@ -540,8 +581,9 @@ class MetaDataset(CachedDataset2):
         :param int sorted_seq_idx:
         :rtype: NumbersDict
         """
-        if self._seq_lens:
-            return self._seq_lens[self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]]
+        if self._seq_lens_file:
+            self._lazy_init_seq_lens()
+            return self._seq_lens[self.get_tag(sorted_seq_idx)]
         return super(MetaDataset, self).get_seq_length(sorted_seq_idx)
 
     def get_tag(self, sorted_seq_idx):
@@ -549,7 +591,10 @@ class MetaDataset(CachedDataset2):
         :param int sorted_seq_idx:
         :rtype: str
         """
-        return self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]
+        if self.seq_list_ordered is not None:
+            return self.seq_list_ordered[self.default_dataset_key][sorted_seq_idx]
+        else:
+            return self.datasets[self.default_dataset_key].get_tag(sorted_seq_idx)
 
     def get_complete_frac(self, sorted_seq_idx: int, **kwargs) -> Optional[float]:
         """
@@ -961,10 +1006,10 @@ class CombinedDataset(CachedDataset2):
         super(CombinedDataset, self).__init__(**kwargs)
         assert self.shuffle_frames_of_nseqs == 0  # not implemented. anyway only for non-recurrent nets
 
+        self.data_map = data_map
         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 +1039,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 +1053,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 +1063,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 +1077,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 +1094,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 +1139,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 +1147,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 +1290,34 @@ 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 not self.datasets[self.dataset_idx2key_map[dataset_idx]].is_less_than_num_seqs(seq_idx):
+            return False
+
+        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 +1340,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], allow_only_lr_suitable=True
+                        )
+                        if self.datasets[self.dataset_idx2key_map[j]].is_less_than_num_seqs(
+                            self.used_num_seqs_per_subset[j]
+                        )
+                        else float("inf"),
+                        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 complete_frac < float("inf"):
+                        break
+                else:
+                    return False  # No dataset has remaining data
+
             elif self.seq_ordering == "random_dataset":
                 while True:
                     # Build probability table
@@ -1323,19 +1424,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
@@ -1348,7 +1453,10 @@ class CombinedDataset(CachedDataset2):
         if dataset_data_key is not None:
             return dataset.get_data(dataset_seq_idx, dataset_data_key)
         else:
-            return numpy.array([], self.data_dtypes[data_key])
+            shape: List[int] = [0] * self.num_outputs[data_key][1]
+            if shape and not self.is_data_sparse(data_key):
+                shape[-1] = self.get_data_dim(data_key)
+            return numpy.zeros(shape, dtype=self.data_dtypes[data_key])
 
     def _collect_single_seq(self, seq_idx):
         """
@@ -1362,19 +1470,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/datasets/multi_proc.py

@@ -42,7 +42,7 @@ class MultiProcDataset(CachedDataset2):
         """
         :param dataset: the dataset to use
         :param num_workers: number of workers to use
-        :param buffer_size: buffer size for each worker, amount of seqs to prefetch
+        :param buffer_size: buffer size for each worker, number of seqs to prefetch
         :param sharding_method: which method to use for sharding the data across the worker procs.
             The default is ``seq_order``, which fetches the full list of seq indices,
             and then distributes shards of that to the other workers.