returnn 1.20240723.142327__tar.gz → 1.20240723.153926__tar.gz

{returnn-1.20240723.142327 → returnn-1.20240723.153926}/PKG-INFO

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

returnn-1.20240723.153926/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240723.153926'
+long_version = '1.20240723.153926+git.7118c54'

{returnn-1.20240723.142327 → returnn-1.20240723.153926}/returnn/datasets/lm.py

@@ -23,6 +23,7 @@ from returnn.log import log
 
 from .basic import DatasetSeq
 from .cached2 import CachedDataset2
+from .util.vocabulary import Vocabulary
 
 
 class LmDataset(CachedDataset2):
@@ -37,6 +38,7 @@ class LmDataset(CachedDataset2):
         self,
         corpus_file,
         skip_empty_lines=True,
+        orth_vocab=None,
         orth_symbols_file=None,
         orth_symbols_map_file=None,
         orth_replace_map_file=None,
@@ -76,6 +78,7 @@ class LmDataset(CachedDataset2):
 
         :param str|()->str|list[str]|()->list[str] corpus_file: Bliss XML or line-based txt. optionally can be gzip.
         :param bool skip_empty_lines: for line-based txt
+        :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>",
                                                        a python dict with {"<symbol>": <index>, ...}
@@ -87,7 +90,7 @@ class LmDataset(CachedDataset2):
         :param str|None seq_end_symbol: what to add at the end, if given.
           will be set as postfix=[seq_end_symbol] or postfix=[] for parse_orth_opts.
         :param str|None unknown_symbol: token to represent unknown words.
-        :param dict[str]|None parse_orth_opts: kwargs for parse_orthography().
+        :param dict[str,typing.Any]|None parse_orth_opts: kwargs for parse_orthography().
         :param dict|None phone_info: A dict containing parameters including a lexicon file for
                                      :class:`LmDataset.PhoneSeqGenerator`.
         :param int add_random_phone_seqs: will add random seqs with the same len as the real seq as additional data.
@@ -102,6 +105,13 @@ class LmDataset(CachedDataset2):
         """
         super(LmDataset, self).__init__(**kwargs)
 
+        self._corpus_file = corpus_file
+        self._skip_empty_lines = skip_empty_lines
+        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
+        self._phone_info = phone_info
+
         if callable(corpus_file):
             corpus_file = corpus_file()
         if callable(orth_symbols_file):
@@ -117,30 +127,31 @@ class LmDataset(CachedDataset2):
         self.word_end_symbol = word_end_symbol
         self.seq_end_symbol = seq_end_symbol
         self.unknown_symbol = unknown_symbol
-        self.parse_orth_opts = parse_orth_opts or {}
-        self.parse_orth_opts.setdefault("word_based", self.word_based)
-        if self.word_end_symbol and not self.word_based:  # Character-based modeling and word_end_symbol is specified.
-            # In this case, sentences end with self.word_end_symbol followed by the self.seq_end_symbol.
-            self.parse_orth_opts.setdefault(
-                "postfix",
-                (
-                    [self.word_end_symbol, self.seq_end_symbol]
-                    if self.seq_end_symbol is not None
-                    else [self.word_end_symbol]
-                ),
-            )
-        else:
-            self.parse_orth_opts.setdefault("postfix", [self.seq_end_symbol] if self.seq_end_symbol is not None else [])
 
-        if orth_symbols_file:
+        self.orth_vocab = None
+        self.orth_symbols = None
+        self.orth_symbols_map = None
+        self.seq_gen = None
+        if orth_vocab:
+            assert not orth_symbols_file, "LmDataset: either orth_vocab or orth_symbols_file"
+            assert not phone_info, "LmDataset: either orth_vocab or phone_info"
+            assert not orth_symbols_map_file, "LmDataset: either orth_vocab or orth_symbols_map_file"
+            assert not auto_replace_unknown_symbol, "LmDataset: auto_replace_unknown_symbol is controlled via the vocab"
+            if isinstance(orth_vocab, dict):
+                self.orth_vocab = Vocabulary.create_vocab(**orth_vocab)
+            elif isinstance(orth_vocab, Vocabulary):
+                self.orth_vocab = orth_vocab
+            else:
+                raise TypeError(f"LmDataset: unexpected orth_vocab type {type(orth_vocab)}")
+            self.labels["data"] = self.orth_vocab.labels
+        elif orth_symbols_file:
             assert not phone_info
             assert not orth_symbols_map_file
             orth_symbols = open(orth_symbols_file).read().splitlines()
             self.orth_symbols_map = {sym: i for (i, sym) in enumerate(orth_symbols)}
             self.orth_symbols = orth_symbols
             self.labels["data"] = orth_symbols
-            self.seq_gen = None
-        if orth_symbols_map_file and orth_symbols_map_file.endswith(".pkl"):
+        elif orth_symbols_map_file and orth_symbols_map_file.endswith(".pkl"):
             import pickle
 
             with open(orth_symbols_map_file, "rb") as f:
@@ -148,7 +159,6 @@ class LmDataset(CachedDataset2):
             self.orth_symbols = self.orth_symbols_map.keys()
             reverse_map = {i: sym for (sym, i) in sorted(self.orth_symbols_map.items())}
             self.labels["data"] = [sym for (i, sym) in sorted(reverse_map.items())]
-            self.seq_gen = None
         elif orth_symbols_map_file:
             assert not phone_info
             with open(orth_symbols_map_file, "r") as f:
@@ -169,30 +179,57 @@ class LmDataset(CachedDataset2):
             self.orth_symbols = [sym for (i, sym) in orth_symbols_imap_list]
             reverse_map = {i: sym for (i, sym) in orth_symbols_imap_list}
             self.labels["data"] = [sym for (i, sym) in sorted(reverse_map.items())]
-            self.seq_gen = None
-        else:
+        elif phone_info is not None:
             assert not orth_symbols_file
             assert isinstance(phone_info, dict)
             self.seq_gen = PhoneSeqGenerator(**phone_info)
-            self.orth_symbols = None
             self.labels["data"] = self.seq_gen.get_class_labels()
-        if orth_replace_map_file:
-            orth_replace_map = load_json(filename=orth_replace_map_file)
-            assert isinstance(orth_replace_map, dict)
-            self.orth_replace_map = {
-                key: parse_orthography_into_symbols(v, word_based=self.word_based)
-                for (key, v) in orth_replace_map.items()
-            }
-            if self.orth_replace_map:
-                if len(self.orth_replace_map) <= 5:
-                    print("  orth_replace_map: %r" % self.orth_replace_map, file=log.v5)
-                else:
-                    print("  orth_replace_map: %i entries" % len(self.orth_replace_map), file=log.v5)
         else:
-            self.orth_replace_map = {}
+            raise ValueError("LmDataset: need orth_symbols_file or orth_symbols_map_file or phone_info")
+
+        self.parse_orth_opts = None
+        if self.orth_symbols is not None:
+            self.parse_orth_opts = parse_orth_opts.copy() if parse_orth_opts else {}
+            self.parse_orth_opts.setdefault("word_based", self.word_based)
+            if self.word_end_symbol and not self.word_based:
+                # Character-based modeling and word_end_symbol is specified.
+                # In this case, sentences end with self.word_end_symbol followed by the self.seq_end_symbol.
+                self.parse_orth_opts.setdefault(
+                    "postfix",
+                    (
+                        [self.word_end_symbol, self.seq_end_symbol]
+                        if self.seq_end_symbol is not None
+                        else [self.word_end_symbol]
+                    ),
+                )
+            else:
+                self.parse_orth_opts.setdefault(
+                    "postfix", [self.seq_end_symbol] if self.seq_end_symbol is not None else []
+                )
+        else:
+            assert not parse_orth_opts
+
+        self.orth_replace_map = None
+        if self.orth_symbols is not None:
+            if orth_replace_map_file:
+                orth_replace_map = load_json(filename=orth_replace_map_file)
+                assert isinstance(orth_replace_map, dict)
+                self.orth_replace_map = {
+                    key: parse_orthography_into_symbols(v, word_based=self.word_based)
+                    for (key, v) in orth_replace_map.items()
+                }
+                if self.orth_replace_map:
+                    if len(self.orth_replace_map) <= 5:
+                        print("  orth_replace_map: %r" % self.orth_replace_map, file=log.v5)
+                    else:
+                        print("  orth_replace_map: %i entries" % len(self.orth_replace_map), file=log.v5)
+            else:
+                self.orth_replace_map = {}
 
-        if word_end_symbol and not word_based:  # Character-based modeling and word_end_symbol is specified.
-            self.orth_replace_map[" "] = [word_end_symbol]  # Replace all spaces by word_end_symbol.
+            if word_end_symbol and not word_based:  # Character-based modeling and word_end_symbol is specified.
+                self.orth_replace_map[" "] = [word_end_symbol]  # Replace all spaces by word_end_symbol.
+        else:
+            assert not orth_replace_map_file
 
         num_labels = len(self.labels["data"])
         if num_labels <= 2**7:
@@ -346,23 +383,10 @@ class LmDataset(CachedDataset2):
             if orth == "</s>":
                 continue  # special sentence end symbol. empty seq, ignore.
 
-            if self.seq_gen:
-                try:
-                    phones = self.seq_gen.generate_seq(orth)
-                except KeyError as e:
-                    if self.log_skipped_seqs:
-                        print(
-                            "LmDataset: skipping sequence %r because of missing lexicon entry: %s" % (orth, e),
-                            file=log.v4,
-                        )
-                        self._reduce_log_skipped_seqs()
-                    if self.error_on_invalid_seq:
-                        raise Exception("LmDataset: invalid seq %r, missing lexicon entry %r" % (orth, e))
-                    self.num_skipped += 1
-                    continue  # try another seq
-                data = self.seq_gen.seq_to_class_idxs(phones, dtype=self.dtype)
+            if self.orth_vocab is not None:
+                data = numpy.array(self.orth_vocab.get_seq(orth), dtype=self.dtype)
 
-            elif self.orth_symbols:
+            elif self.orth_symbols is not None:
                 orth_syms = parse_orthography(orth, **self.parse_orth_opts)
                 while True:
                     orth_syms = sum([self.orth_replace_map.get(s, [s]) for s in orth_syms], [])
@@ -416,8 +440,24 @@ class LmDataset(CachedDataset2):
                     self.num_skipped += 1
                     continue  # try another seq
 
+            elif self.seq_gen is not None:
+                try:
+                    phones = self.seq_gen.generate_seq(orth)
+                except KeyError as e:
+                    if self.log_skipped_seqs:
+                        print(
+                            "LmDataset: skipping sequence %r because of missing lexicon entry: %s" % (orth, e),
+                            file=log.v4,
+                        )
+                        self._reduce_log_skipped_seqs()
+                    if self.error_on_invalid_seq:
+                        raise Exception("LmDataset: invalid seq %r, missing lexicon entry %r" % (orth, e))
+                    self.num_skipped += 1
+                    continue  # try another seq
+                data = self.seq_gen.seq_to_class_idxs(phones, dtype=self.dtype)
+
             else:
-                assert False
+                assert False, f"{self}, {self.orth_vocab}, {self.orth_symbols}, {self.seq_gen}"
 
             targets = {}
             for i in range(self.add_random_phone_seqs):

{returnn-1.20240723.142327 → returnn-1.20240723.153926}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240723.142327 → returnn-1.20240723.153926}/tests/test_Dataset.py

@@ -584,6 +584,158 @@ def test_OggZipDataset():
         assert classes_ == _demo_txt + "."
 
 
+def test_LmDataset_char_based():
+    from returnn.datasets.lm import LmDataset
+
+    with tempfile.NamedTemporaryFile("wt", suffix=".txt") as txt_file, tempfile.NamedTemporaryFile(
+        "wt", suffix=".syms"
+    ) as orth_syms_file:
+        txt_file.write("Hello world\n")
+        txt_file.write("Next line\n")
+        txt_file.flush()
+        orth_syms_file.write("[END]\n")
+        covered = set()
+        for c in open(txt_file.name, "r").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": txt_file.name,
+                "orth_symbols_file": orth_syms_file.name,
+            }
+        )
+        assert isinstance(dataset, LmDataset)
+        dataset.init_seq_order(epoch=1)
+        dataset.load_seqs(0, 2)
+        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]
+        assert not dataset.is_less_than_num_seqs(2)
+
+
+def test_LmDataset_word_based():
+    from returnn.datasets.lm import LmDataset
+
+    with tempfile.NamedTemporaryFile("wt", suffix=".txt") as txt_file, tempfile.NamedTemporaryFile(
+        "wt", suffix=".syms"
+    ) as orth_syms_file:
+        txt_file.write("Hello world\n")
+        txt_file.write("Next line\n")
+        txt_file.flush()
+        orth_syms_file.write("[END]\n")
+        orth_syms_file.write("Hello\n")
+        orth_syms_file.write("world\n")
+        orth_syms_file.write("Next\n")
+        orth_syms_file.write("line\n")
+        orth_syms_file.flush()
+
+        dataset = init_dataset(
+            {
+                "class": "LmDataset",
+                "corpus_file": txt_file.name,
+                "word_based": True,
+                "orth_symbols_file": orth_syms_file.name,
+            }
+        )
+        assert isinstance(dataset, LmDataset)
+        dataset.init_seq_order(epoch=1)
+        dataset.load_seqs(0, 2)
+        orth = dataset.get_data(0, "data")
+        assert orth.tolist() == [1, 2, 0]
+        orth = dataset.get_data(1, "data")
+        assert orth.tolist() == [3, 4, 0]
+        assert not dataset.is_less_than_num_seqs(2)
+
+
+def test_LmDataset_vocab_based():
+    from returnn.datasets.lm import LmDataset
+
+    with tempfile.NamedTemporaryFile("wt", suffix=".txt") as txt_file, tempfile.NamedTemporaryFile(
+        "wt", suffix=".syms"
+    ) as orth_syms_file:
+        txt_file.write("Hello world\n")
+        txt_file.write("Next line\n")
+        txt_file.flush()
+        orth_syms_file.write(
+            repr(
+                {
+                    "[END]": 0,
+                    "Hello": 1,
+                    "world": 2,
+                    "Next": 3,
+                    "line": 4,
+                }
+            )
+        )
+        orth_syms_file.write("\n")
+        orth_syms_file.flush()
+
+        dataset = init_dataset(
+            {
+                "class": "LmDataset",
+                "corpus_file": txt_file.name,
+                "orth_vocab": {
+                    "class": "SamplingBytePairEncoding",  # just as an example
+                    "vocab_file": orth_syms_file.name,
+                    "breadth_prob": 0.0,
+                    "unknown_label": None,
+                },
+            }
+        )
+        assert isinstance(dataset, LmDataset)
+        dataset.init_seq_order(epoch=1)
+        dataset.load_seqs(0, 2)
+        orth = dataset.get_data(0, "data")
+        assert orth.tolist() == [1, 2]
+        orth = dataset.get_data(1, "data")
+        assert orth.tolist() == [3, 4]
+        assert not dataset.is_less_than_num_seqs(2)
+
+
+def test_LmDataset_pickle():
+    import pickle
+    from returnn.datasets.lm import LmDataset
+
+    with tempfile.NamedTemporaryFile("wt", suffix=".txt") as txt_file, tempfile.NamedTemporaryFile(
+        "wt", suffix=".syms"
+    ) as orth_syms_file:
+        txt_file.write("Hello world\n")
+        txt_file.write("Next line\n")
+        txt_file.flush()
+        orth_syms_file.write(repr({"[END]": 0, "Hello": 1, "world": 2, "Next": 3, "line": 4}))
+        orth_syms_file.write("\n")
+        orth_syms_file.flush()
+
+        dataset = init_dataset(
+            {
+                "class": "LmDataset",
+                "corpus_file": txt_file.name,
+                "orth_vocab": {
+                    "vocab_file": orth_syms_file.name,
+                    "unknown_label": None,
+                },
+            }
+        )
+        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, 2)
+        orth = dataset.get_data(0, "data")
+        assert orth.tolist() == [1, 2]
+        orth = dataset.get_data(1, "data")
+        assert orth.tolist() == [3, 4]
+        assert not dataset.is_less_than_num_seqs(2)
+
+
 def test_MetaDataset():
     _demo_txt = "some utterance text"
 

returnn-1.20240723.142327/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240723.142327'
-long_version = '1.20240723.142327+git.f1f1c76'