returnn 1.20241210.111636__tar.gz → 1.20241211.145755__tar.gz

{returnn-1.20241210.111636 → returnn-1.20241211.145755}/PKG-INFO

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

returnn-1.20241211.145755/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241211.145755'
+long_version = '1.20241211.145755+git.4c6e245'

{returnn-1.20241210.111636 → returnn-1.20241211.145755}/returnn/datasets/audio.py

@@ -167,6 +167,7 @@ class OggZipDataset(CachedDataset2):
         else:
             assert isinstance(epoch_wise_filter, EpochWiseFilter)
             self.epoch_wise_filter = epoch_wise_filter
+        self._seq_index_by_tag: Optional[Dict[str, int]] = None
         self._seq_order = None  # type: typing.Optional[typing.Sequence[int]]
 
     def _read(self, filename, zip_index):
@@ -288,6 +289,11 @@ class OggZipDataset(CachedDataset2):
             segment_file_handle = open(segment_file)
             self.segments = set(segment_file_handle.read().splitlines())
 
+    def _lazy_init_seq_index_by_tag(self):
+        if self._seq_index_by_tag is not None:
+            return
+        self._seq_index_by_tag = {self._get_tag_from_info_dict(seq): i for i, seq in enumerate(self._data)}
+
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """
         If random_shuffle_epoch1, for epoch 1 with "random" ordering, we leave the given order as is.
@@ -327,12 +333,12 @@ class OggZipDataset(CachedDataset2):
         if seq_order is not None:
             self._seq_order = seq_order
         elif seq_list is not None:
-            seqs = {self._get_tag_from_info_dict(seq): i for i, seq in enumerate(self._data)}
+            self._lazy_init_seq_index_by_tag()
             for seq_tag in seq_list:
-                assert seq_tag in seqs, "Did not find all requested seqs. We have eg: %s" % (
+                assert seq_tag in self._seq_index_by_tag, "Did not find all requested seqs. We have eg: %s" % (
                     self._get_tag_from_info_dict(self._data[0]),
                 )
-            self._seq_order = [seqs[seq_tag] for seq_tag in seq_list]
+            self._seq_order = [self._seq_index_by_tag[seq_tag] for seq_tag in seq_list]
         else:
             num_seqs = len(self._data)
             self._seq_order = self.get_seq_order_for_epoch(epoch=epoch, num_seqs=num_seqs, get_seq_len=get_seq_len)

{returnn-1.20241210.111636 → returnn-1.20241211.145755}/returnn/datasets/basic.py

@@ -330,6 +330,8 @@ class Dataset:
             assert chunk_step.max_value() > 0, "chunking step must be positive (for some key)"
         return chunk_size, chunk_step, None
 
+    _seq_list_file_cache: Dict[str, Union[List[str], Dict[str, List[str]]]] = {}
+
     @staticmethod
     def _load_seq_list_file(
         filename: str, *, use_cache_manager: bool = False, expect_list: bool = True
@@ -340,6 +342,12 @@ class Dataset:
         :param expect_list:
         :return: seq list file content. usually a list of seqs.
         """
+        if filename in Dataset._seq_list_file_cache:
+            seq_list = Dataset._seq_list_file_cache[filename]
+            if expect_list:
+                assert isinstance(seq_list, list)
+            return seq_list
+        _filename = filename
         if use_cache_manager:
             import returnn.util.basic
 
@@ -367,6 +375,7 @@ class Dataset:
                 assert isinstance(seq_list, list)
         else:
             seq_list = _open("rt").read().splitlines()
+        Dataset._seq_list_file_cache[_filename] = seq_list
         return seq_list
 
     def _sliding_window(self, xr):

{returnn-1.20241210.111636 → returnn-1.20241211.145755}/returnn/datasets/lm.py

@@ -46,8 +46,10 @@ class LmDataset(CachedDataset2):
     def __init__(
         self,
         corpus_file,
+        *,
         use_cache_manager=False,
         skip_empty_lines=True,
+        seq_list_file=None,
         orth_vocab=None,
         orth_symbols_file=None,
         orth_symbols_map_file=None,
@@ -89,6 +91,8 @@ class LmDataset(CachedDataset2):
         :param str|()->str|list[str]|()->list[str] corpus_file: Bliss XML or line-based txt. optionally can be gzip.
         :param bool use_cache_manager: uses :func:`returnn.util.basic.cf`
         :param bool skip_empty_lines: for line-based txt
+        :param str|list[str]|None seq_list_file: optional custom seq tags to use instead of the "line-%i" seq tags.
+            Pickle (.pkl) or txt (line-based seq tags). Optionally gzipped (.gz).
         :param dict[str,typing.Any]|Vocabulary orth_vocab:
         :param str|()->str|None orth_symbols_file: a text file containing a list of orthography symbols
         :param str|()->str|None orth_symbols_map_file: either a list of orth symbols, each line: "<symbol> <index>",
@@ -119,6 +123,7 @@ class LmDataset(CachedDataset2):
         self._corpus_file = corpus_file
         self._use_cache_manager = use_cache_manager
         self._skip_empty_lines = skip_empty_lines
+        self._seq_list_file = seq_list_file
         self._orth_symbols_file = orth_symbols_file
         self._orth_symbols_map_file = orth_symbols_map_file
         self._orth_replace_map_file = orth_replace_map_file
@@ -274,6 +279,8 @@ class LmDataset(CachedDataset2):
         self._orth_files: Optional[List[BinaryIO]] = None
         self._orth_mmaps = None
         self._orths_offsets_and_lens: Optional[List[Tuple[int, int]]] = None  # will be loaded in _lazy_init
+        self._seq_list: Optional[List[str]] = None
+        self._seq_index_by_tag: Optional[dict[str, int]] = None
 
         self.next_orth_idx = 0
         self.next_seq_idx = 0
@@ -301,6 +308,7 @@ class LmDataset(CachedDataset2):
         total_bytes_read = 0
         orths = []
         self._orths_offsets_and_lens = orths
+        lens_per_corpus_file = []
         start_time = time.time()
         last_print_time = start_time
 
@@ -341,6 +349,7 @@ class LmDataset(CachedDataset2):
         # If a list of files is provided, concatenate all.
         if not isinstance(corpus_file, list):
             corpus_file = [corpus_file]
+        prev_orth_len = 0
         for file_name in corpus_file:
             if self._use_cache_manager:
                 file_name = cf(file_name)
@@ -376,10 +385,31 @@ class LmDataset(CachedDataset2):
                     pos = next_new_line + 1
                     _maybe_report_status()
 
+            lens_per_corpus_file.append(len(orths) - prev_orth_len)
+            prev_orth_len = len(orths)
+
         if tmp_file is not None:
             tmp_file.flush()
             self._orth_mmaps[tmp_file_orth_files_index] = mmap.mmap(tmp_file.fileno(), 0, flags=mmap.MAP_PRIVATE)
 
+        if self._seq_list_file:
+            if isinstance(self._seq_list_file, str):
+                seq_list: List[str] = self._load_seq_list_file(
+                    self._seq_list_file, use_cache_manager=self._use_cache_manager
+                )
+            elif isinstance(self._seq_list_file, list):
+                assert len(self._seq_list_file) == len(lens_per_corpus_file)
+                seq_list: List[str] = []
+                for i, fn in enumerate(self._seq_list_file):
+                    seq_list_: List[str] = self._load_seq_list_file(fn, use_cache_manager=self._use_cache_manager)
+                    assert len(seq_list_) == lens_per_corpus_file[i]
+                    seq_list.extend(seq_list_)
+            else:
+                raise TypeError(f"invalid seq_list_file type {type(self._seq_list_file).__name__}")
+            assert isinstance(seq_list, list)
+            assert len(self._orths_offsets_and_lens) == len(seq_list)
+            self._seq_list = seq_list
+
         print(
             f"  done, loaded {len(self._orths_offsets_and_lens)} sequences,"
             f" {human_bytes_size(total_bytes_read)},"
@@ -445,8 +475,16 @@ class LmDataset(CachedDataset2):
         if seq_order is not None:
             self.seq_order = seq_order
         elif seq_list is not None:
-            assert all(s.startswith(self._tag_prefix) for s in seq_list)
-            self.seq_order = [int(s[len(self._tag_prefix) :]) for s in seq_list]
+            # Might not be initialized. Can even do without init. Thus check seq_list_file.
+            if self._seq_list_file is None:
+                assert all(s.startswith(self._tag_prefix) for s in seq_list)
+                self.seq_order = [int(s[len(self._tag_prefix) :]) for s in seq_list]
+            else:
+                # Need seq list for this. Just do the lazy init now.
+                self._lazy_init()
+                if self._seq_index_by_tag is None:
+                    self._seq_index_by_tag = {tag: i for (i, tag) in enumerate(self._seq_list)}
+                self.seq_order = [self._seq_index_by_tag[s] for s in seq_list]
         elif epoch is None:
             self.seq_order = []
         else:
@@ -486,6 +524,9 @@ class LmDataset(CachedDataset2):
 
     def get_all_tags(self) -> List[str]:
         """:return: all seq tags"""
+        self._lazy_init()
+        if self._seq_list is not None:
+            return self._seq_list
         num_seqs = self.get_total_num_seqs()
         return [self._tag_prefix + str(line_nr) for line_nr in range(num_seqs)]
 
@@ -525,10 +566,11 @@ class LmDataset(CachedDataset2):
             # get sequence for the next index given by seq_order
             idx, offset, len_ = self._orths_offsets_and_lens[true_idx]
             orth = self._orth_mmaps[idx][offset : offset + len_].decode("utf8").strip()
-            seq_tag = self._tag_prefix + str(true_idx)
+            if self._seq_list is None:
+                seq_tag = self._tag_prefix + str(true_idx)
+            else:
+                seq_tag = self._seq_list[true_idx]
             self.next_orth_idx += 1
-            if orth == "</s>":
-                continue  # special sentence end symbol. empty seq, ignore.
 
             if self.orth_vocab is not None:
                 data = numpy.array(self.orth_vocab.get_seq(orth), dtype=self.dtype)

{returnn-1.20241210.111636 → returnn-1.20241211.145755}/returnn.egg-info/PKG-INFO

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

{returnn-1.20241210.111636 → returnn-1.20241211.145755}/tests/test_Config.py

@@ -1,7 +1,6 @@
 import sys
 import _setup_test_env  # noqa
 import unittest
-from nose.tools import assert_equal, assert_is_instance, assert_in, assert_greater, assert_true, assert_false
 from pprint import pprint
 from returnn.config import *
 from returnn.util import better_exchook
@@ -21,16 +20,16 @@ def test_old_format():
         )
     )
 
-    assert_true(config.has("num_inputs"))
-    assert_true(config.has("hidden_type"))
-    assert_equal(config.int("num_inputs", -1), 3)
-    assert_equal(config.value("hidden_type", "x"), "forward,lstm")
-    assert_equal(config.value("hidden_type", "x", index=0), "forward")
-    assert_equal(config.value("hidden_type", "x", index=1), "lstm")
-    assert_equal(config.list("hidden_type", ["x"]), ["forward", "lstm"])
+    assert config.has("num_inputs") is True
+    assert config.has("hidden_type") is True
+    assert config.int("num_inputs", -1) == 3
+    assert config.value("hidden_type", "x") == "forward,lstm"
+    assert config.value("hidden_type", "x", index=0) == "forward"
+    assert config.value("hidden_type", "x", index=1) == "lstm"
+    assert config.list("hidden_type", ["x"]) == ["forward", "lstm"]
 
-    assert_false(config.is_typed("num_inputs"))
-    assert_false(config.is_typed("hidden_type"))
+    assert config.is_typed("num_inputs") is False
+    assert config.is_typed("hidden_type") is False
 
 
 def test_json_format():
@@ -47,19 +46,19 @@ def test_json_format():
         )
     )
 
-    assert_true(config.has("num_inputs"))
-    assert_true(config.has("hidden_type"))
-    assert_equal(config.int("num_inputs", -1), 3)
-    assert_equal(config.value("hidden_type", "x"), "forward,lstm")
-    assert_equal(config.value("hidden_type", "x", index=0), "forward")
-    assert_equal(config.value("hidden_type", "x", index=1), "lstm")
-    assert_equal(config.list("hidden_type", ["x"]), ["forward", "lstm"])
+    assert config.has("num_inputs") is True
+    assert config.has("hidden_type") is True
+    assert config.int("num_inputs", -1) == 3
+    assert config.value("hidden_type", "x") == "forward,lstm"
+    assert config.value("hidden_type", "x", index=0) == "forward"
+    assert config.value("hidden_type", "x", index=1) == "lstm"
+    assert config.list("hidden_type", ["x"]) == ["forward", "lstm"]
 
-    assert_true(config.is_typed("num_inputs"))
-    assert_true(config.is_typed("hidden_type"))
-    assert_is_instance(config.typed_value("num_inputs"), int)
-    assert_is_instance(config.typed_value("hidden_type"), list)
-    assert_equal(config.typed_value("hidden_type"), ["forward", "lstm"])
+    assert config.is_typed("num_inputs") is True
+    assert config.is_typed("hidden_type") is True
+    assert isinstance(config.typed_value("num_inputs"), int)
+    assert isinstance(config.typed_value("hidden_type"), list)
+    assert config.typed_value("hidden_type") == ["forward", "lstm"]
 
 
 def test_py_config():
@@ -74,19 +73,19 @@ hidden_type = ["forward", "lstm"]
         )
     )
 
-    assert_true(config.has("num_inputs"))
-    assert_true(config.has("hidden_type"))
-    assert_equal(config.int("num_inputs", -1), 3)
-    assert_equal(config.value("hidden_type", "x"), "forward,lstm")
-    assert_equal(config.value("hidden_type", "x", index=0), "forward")
-    assert_equal(config.value("hidden_type", "x", index=1), "lstm")
-    assert_equal(config.list("hidden_type", ["x"]), ["forward", "lstm"])
+    assert config.has("num_inputs") is True
+    assert config.has("hidden_type") is True
+    assert config.int("num_inputs", -1) == 3
+    assert config.value("hidden_type", "x") == "forward,lstm"
+    assert config.value("hidden_type", "x", index=0) == "forward"
+    assert config.value("hidden_type", "x", index=1) == "lstm"
+    assert config.list("hidden_type", ["x"]) == ["forward", "lstm"]
 
-    assert_true(config.is_typed("num_inputs"))
-    assert_true(config.is_typed("hidden_type"))
-    assert_is_instance(config.typed_value("num_inputs"), int)
-    assert_is_instance(config.typed_value("hidden_type"), list)
-    assert_equal(config.typed_value("hidden_type"), ["forward", "lstm"])
+    assert config.is_typed("num_inputs") is True
+    assert config.is_typed("hidden_type") is True
+    assert isinstance(config.typed_value("num_inputs"), int)
+    assert isinstance(config.typed_value("hidden_type"), list)
+    assert config.typed_value("hidden_type") == ["forward", "lstm"]
 
 
 def test_rnn_init_config_py_global_var():
@@ -115,9 +114,9 @@ def test_func():
     assert rnn.config.has("task")
     assert rnn.config.has("test_value")
     assert rnn.config.has("test_func")
-    assert_equal(rnn.config.value("task", None), "search")
+    assert rnn.config.value("task", None) == "search"
     assert rnn.config.is_typed("test_value")
-    assert_equal(rnn.config.typed_value("test_value"), 42)
+    assert rnn.config.typed_value("test_value") == 42
     assert rnn.config.is_typed("test_func")
     # So far it's fine.
     # Now something a bit strange.
@@ -125,7 +124,7 @@ def test_func():
     assert rnn.config.is_typed("task")
     test_func = rnn.config.typed_dict["test_func"]
     assert callable(test_func)
-    assert_equal(test_func(), "search")
+    assert test_func() == "search"
 
 
 def test_rnn_init_config_py_cmd_type():
@@ -152,7 +151,7 @@ def test_func():
     assert rnn.config.is_typed("test_func")
     test_func = rnn.config.typed_dict["test_func"]
     assert callable(test_func)
-    assert_equal(test_func(), 0)
+    assert test_func() == 0
 
 
 def test_config_py_ext():
@@ -176,9 +175,9 @@ def test_func():
     assert config.is_typed("test_func")
     test_func = config.typed_dict["test_func"]
     assert callable(test_func)
-    assert_equal(test_func(), "train")
+    assert test_func() == "train"
     config.set("task", "search")
-    assert_equal(test_func(), "search")
+    assert test_func() == "search"
 
 
 def test_config_py_old_returnn_imports():