returnn 1.20240725.5736__tar.gz → 1.20240725.135013__tar.gz

{returnn-1.20240725.5736 → returnn-1.20240725.135013}/PKG-INFO

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

returnn-1.20240725.135013/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240725.135013'
+long_version = '1.20240725.135013+git.31eea55'

{returnn-1.20240725.5736 → returnn-1.20240725.135013}/returnn/__main__.py

@@ -536,9 +536,9 @@ def execute_main_task():
     if config.is_true("dry_run"):
         print("Dry run, will not save anything.", file=log.v1)
     if task == "train":
-        assert (
-            train_data and train_data.have_seqs()
-        ), "no train files specified, check 'train' option: %s" % config.value("train", None)
+        # Avoid too many investigations on train_data (like have_seqs()),
+        # to avoid triggering any lazy init, which could be unnecessary in the main proc.
+        assert train_data, "no train files specified, check 'train' option: %s" % config.value("train", None)
         engine.init_train_from_config(config, train_data, dev_data, eval_data)
         engine.train()
     elif task == "eval":

{returnn-1.20240725.5736 → returnn-1.20240725.135013}/returnn/datasets/lm.py

@@ -7,7 +7,7 @@ and some related helpers.
 
 from __future__ import annotations
 
-from typing import Optional, Union, Callable, List, Tuple
+from typing import Optional, Union, Callable, List, Tuple, BinaryIO, cast
 import typing
 import os
 import sys
@@ -193,7 +193,7 @@ class LmDataset(CachedDataset2):
             self.seq_gen = PhoneSeqGenerator(**phone_info)
             self.labels["data"] = self.seq_gen.get_class_labels()
         else:
-            raise ValueError("LmDataset: need orth_symbols_file or orth_symbols_map_file or phone_info")
+            raise ValueError("LmDataset: need orth_vocab or orth_symbols_file or orth_symbols_map_file or phone_info")
 
         self.parse_orth_opts = None
         if self.orth_symbols is not None:
@@ -271,8 +271,8 @@ class LmDataset(CachedDataset2):
             self.num_outputs["delayed"] = self.num_outputs["data"]
             self.labels["delayed"] = self.labels["data"]
 
-        self._orth_tmp_file = None
-        self._orth_mmap = None
+        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.next_orth_idx = 0
@@ -293,46 +293,96 @@ class LmDataset(CachedDataset2):
         import tempfile
         import mmap
 
-        tmp_file = tempfile.NamedTemporaryFile(prefix="returnn_lm_dataset_", suffix="_tmp.txt")
+        self._orth_mmaps = []
+        self._orth_files = []
+        tmp_file_orth_files_index: Optional[int] = None
+        tmp_file: Optional[BinaryIO] = None
+        tmp_file_offset = 0
+        total_bytes_read = 0
         orths = []
-        offset = 0
-        self._orth_tmp_file = tmp_file
         self._orths_offsets_and_lens = orths
         start_time = time.time()
         last_print_time = start_time
 
-        def _add_line(line: bytes):
-            nonlocal offset, last_print_time
+        def _init_tmp_file():
+            nonlocal tmp_file, tmp_file_orth_files_index
+
+            if tmp_file is not None:
+                return
+            tmp_file = tempfile.NamedTemporaryFile(prefix="returnn_lm_dataset_", suffix="_tmp.txt")
+            tmp_file = cast(BinaryIO, tmp_file)
+            tmp_file_orth_files_index = len(self._orth_files)
+            self._orth_files.append(tmp_file)
+            self._orth_mmaps.append(None)  # will be set later
 
-            orths.append((offset, len(line)))
+        def _tmp_file_add_line(line: bytes):
+            nonlocal tmp_file_offset, total_bytes_read
+
+            orths.append((tmp_file_orth_files_index, tmp_file_offset, len(line)))
             tmp_file.write(line)
             tmp_file.write(b"\n")
-            offset += len(line) + 1
+            tmp_file_offset += len(line) + 1
+            total_bytes_read += len(line) + 1
+
+            _maybe_report_status()
+
+        def _maybe_report_status():
+            nonlocal last_print_time
 
             if time.time() - last_print_time > 10:
                 print(
-                    f"  ..., loaded {len(self._orths_offsets_and_lens)} sequences,"
-                    f" {human_bytes_size(offset)},"
+                    f"  ... loaded {len(self._orths_offsets_and_lens)} sequences,"
+                    f" {human_bytes_size(total_bytes_read)},"
                     f" after {hms(time.time() - start_time)}",
                     file=log.v4,
                 )
                 last_print_time = time.time()
 
-        if isinstance(corpus_file, list):  # If a list of files is provided, concatenate all.
-            for file_name in corpus_file:
-                if self._use_cache_manager:
-                    file_name = cf(file_name)
-                iter_corpus(file_name, skip_empty_lines=self._skip_empty_lines, decode=False, callback=_add_line)
-        else:
+        # If a list of files is provided, concatenate all.
+        if not isinstance(corpus_file, list):
+            corpus_file = [corpus_file]
+        for file_name in corpus_file:
             if self._use_cache_manager:
-                corpus_file = cf(corpus_file)
-            iter_corpus(corpus_file, skip_empty_lines=self._skip_empty_lines, decode=False, callback=_add_line)
+                file_name = cf(file_name)
+            if _is_bliss(file_name):
+                _init_tmp_file()
+                _iter_bliss(filename=file_name, callback=_tmp_file_add_line, decode=False)
+            elif file_name.endswith(".gz"):
+                _init_tmp_file()
+                _iter_txt(
+                    filename=file_name,
+                    callback=_tmp_file_add_line,
+                    skip_empty_lines=self._skip_empty_lines,
+                    decode=False,
+                )
+            else:  # Raw txt file
+                # Directly mmap the file.
+                # We just need to scan once through it to find line offsets.
+                file = open(file_name, "rb")
+                file_mmap = mmap.mmap(file.fileno(), 0, flags=mmap.MAP_PRIVATE)
+                file_index = len(self._orth_files)
+                self._orth_files.append(file)
+                self._orth_mmaps.append(file_mmap)
+
+                pos = 0
+                while True:
+                    next_new_line = file_mmap.find(b"\n", pos)
+                    if next_new_line == -1:
+                        break
+                    line_len = next_new_line - pos
+                    if line_len or not self._skip_empty_lines:
+                        orths.append((file_index, pos, line_len))
+                    total_bytes_read += line_len + 1
+                    pos = next_new_line + 1
+                    _maybe_report_status()
+
+        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)
 
-        tmp_file.flush()
-        self._orth_mmap = mmap.mmap(tmp_file.fileno(), 0, flags=mmap.MAP_PRIVATE)
         print(
             f"  done, loaded {len(self._orths_offsets_and_lens)} sequences,"
-            f" {human_bytes_size(offset)},"
+            f" {human_bytes_size(total_bytes_read)},"
             f" in {hms(time.time() - start_time)}",
             file=log.v4,
         )
@@ -449,8 +499,8 @@ class LmDataset(CachedDataset2):
             true_idx = self.seq_order[self.next_orth_idx]
             self._lazy_init()
             # get sequence for the next index given by seq_order
-            offset, len_ = self._orths_offsets_and_lens[true_idx]
-            orth = self._orth_mmap[offset : offset + len_].decode("utf8")
+            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)
             self.next_orth_idx += 1
             if orth == "</s>":

{returnn-1.20240725.5736 → returnn-1.20240725.135013}/returnn/torch/data/returnn_dataset_wrapper.py

@@ -5,11 +5,14 @@ We make use of torch.utils.data.IterDataPipe.
 """
 
 from __future__ import annotations
+
 from typing import Callable, Optional, Iterable, Dict
+import sys
 import numpy
 import torch.utils.data
 from returnn.datasets.basic import Dataset as ReturnnDataset
 from returnn.datasets.util.strings import str_to_numpy_array
+from returnn.util.better_exchook import better_exchook
 
 ResetCallbackT = Callable[[], None]
 
@@ -72,15 +75,22 @@ class ReturnnDatasetIterDataPipe(torch.utils.data.IterDataPipe):
         """
         :return: generator providing data samples in the form of a dict data_key -> data
         """
-        data_keys = self._dataset.get_data_keys()
-
-        seq_index = 0
-        while self._dataset.is_less_than_num_seqs(seq_index):
-            self._dataset.load_seqs(seq_index, seq_index + 1)
-            data = {data_key: self._dataset.get_data(seq_index, data_key) for data_key in data_keys}
-            data["seq_tag"] = str_to_numpy_array(self._dataset.get_tag(seq_index))
-            yield data
-            seq_index += 1
+        try:
+            data_keys = self._dataset.get_data_keys()
+
+            seq_index = 0
+            while self._dataset.is_less_than_num_seqs(seq_index):
+                self._dataset.load_seqs(seq_index, seq_index + 1)
+                data = {data_key: self._dataset.get_data(seq_index, data_key) for data_key in data_keys}
+                data["seq_tag"] = str_to_numpy_array(self._dataset.get_tag(seq_index))
+                yield data
+                seq_index += 1
+
+        except Exception:
+            # Report here as well, because the parent caller (PyTorch DataLoader) might first catch it
+            # and only report a stripped version.
+            better_exchook(*sys.exc_info(), autodebugshell=False)
+            raise
 
     def __getitem__(self, index):
         raise Exception(f"{self.__class__.__name__}.__getitem__ not supported")

{returnn-1.20240725.5736 → returnn-1.20240725.135013}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240725.5736 → returnn-1.20240725.135013}/tests/test_Dataset.py

@@ -739,6 +739,56 @@ def test_LmDataset_pickle():
         assert not dataset.is_less_than_num_seqs(2)
 
 
+def test_LmDataset_multiple_and_gzipped():
+    import pickle
+    import gzip
+    from returnn.datasets.lm import LmDataset
+
+    with tempfile.NamedTemporaryFile("wb", suffix=".txt.gz") as gz_txt_file, tempfile.NamedTemporaryFile(
+        "wb", suffix=".txt.gz"
+    ) as gz_txt_file2, tempfile.NamedTemporaryFile("wt", suffix=".txt") as txt_file3, tempfile.NamedTemporaryFile(
+        "wt", suffix=".syms"
+    ) as orth_syms_file:
+        with gzip.GzipFile(gz_txt_file.name, "wb") as txt_file:
+            txt_file.write(b"Hello world\n")
+            txt_file.write(b"Next line\n")
+        with gzip.GzipFile(gz_txt_file2.name, "wb") as txt_file2:
+            txt_file2.write(b"Next world\n")
+        txt_file3.write("Hello line\n")
+        txt_file3.flush()
+        orth_syms_file.write("[END]\n")
+        covered = set()
+        for c in gzip.open(gz_txt_file.name, "rt").read():
+            if c not in covered and c != "\n":
+                orth_syms_file.write(c + "\n")
+                covered.add(c)
+        orth_syms_file.flush()
+
+        dataset = init_dataset(
+            {
+                "class": "LmDataset",
+                "corpus_file": [gz_txt_file.name, gz_txt_file2.name, txt_file3.name],
+                "orth_symbols_file": orth_syms_file.name,
+            }
+        )
+        assert isinstance(dataset, LmDataset)
+        s = pickle.dumps(dataset)
+        dataset = pickle.loads(s)
+        assert isinstance(dataset, LmDataset)
+
+        dataset.init_seq_order(epoch=1)
+        dataset.load_seqs(0, 4)
+        orth = dataset.get_data(0, "data")
+        assert orth.tolist() == [1, 2, 3, 3, 4, 5, 6, 4, 7, 3, 8, 0]
+        orth = dataset.get_data(1, "data")
+        assert orth.tolist() == [9, 2, 10, 11, 5, 3, 12, 13, 2, 0]
+        orth = dataset.get_data(2, "data")
+        assert orth.tolist() == [9, 2, 10, 11, 5, 6, 4, 7, 3, 8, 0]
+        orth = dataset.get_data(3, "data")
+        assert orth.tolist() == [1, 2, 3, 3, 4, 5, 3, 12, 13, 2, 0]
+        assert not dataset.is_less_than_num_seqs(4)
+
+
 def test_MetaDataset():
     _demo_txt = "some utterance text"
 

returnn-1.20240725.5736/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240725.005736'
-long_version = '1.20240725.005736+git.d5855be'