returnn 1.20240620.64807__tar.gz → 1.20240620.201955__tar.gz

{returnn-1.20240620.64807/returnn.egg-info → returnn-1.20240620.201955}/PKG-INFO

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

returnn-1.20240620.201955/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240620.201955'
+long_version = '1.20240620.201955+git.7a0206a'

{returnn-1.20240620.64807 → returnn-1.20240620.201955}/returnn/datasets/util/vocabulary.py

@@ -19,6 +19,7 @@ import sys
 import numpy
 
 from returnn.log import log
+from returnn.util.basic import NotSpecified
 
 
 class Vocabulary(object):
@@ -327,14 +328,48 @@ class SamplingBytePairEncoding(Vocabulary):
     This will encode the text on-the-fly with BPE.
     """
 
-    def __init__(self, vocab_file, breadth_prob, seq_postfix=None, **kwargs):
+    def __init__(
+        self,
+        vocab_file: str,
+        breadth_prob: float,
+        seq_postfix: Optional[List[int]] = None,
+        label_postfix_merge_symbol: Optional[str] = NotSpecified,
+        word_prefix_symbol: Optional[str] = NotSpecified,
+        **kwargs,
+    ):
         """
-        :param str vocab_file:
-        :param float breadth_prob:
-        :param list[int]|None seq_postfix: labels will be added to the seq in self.get_seq
+        :param vocab_file:
+        :param breadth_prob:
+        :param seq_postfix: labels will be added to the seq in self.get_seq
+        :param label_postfix_merge_symbol: If given, will use this as label postfix merge symbol,
+            i.e. when this occurs at the end of a label, it is supposed to be merged with the next label,
+            i.e. the space between them is removed and is not a word boundary.
+            If None, will not use any postfix merge symbol.
+            If not specified, and also word_prefix_symbol is not specified, will use "@@" by default here,
+            the standard from subword-nmt, and our original behavior.
+        :param word_prefix_symbol: If given, every new word starts with this symbol.
+            This also implies that there are no spaces between words
+            and this symbol is a placeholder for the space.
+            If None, will not use this logic.
+            For SentencePiece, you usually would use "▁" here.
         """
         super(SamplingBytePairEncoding, self).__init__(vocab_file=vocab_file, seq_postfix=seq_postfix, **kwargs)
-        from returnn.util.bpe import SamplingBytePairEncoder
+        from returnn.util.bpe import SamplingBytePairEncoder, BpePostMergeSymbol, BpeOpts
+
+        if label_postfix_merge_symbol is NotSpecified and word_prefix_symbol is NotSpecified:
+            label_postfix_merge_symbol = BpePostMergeSymbol
+            word_prefix_symbol = None
+        else:
+            if label_postfix_merge_symbol is NotSpecified:
+                label_postfix_merge_symbol = None
+            if word_prefix_symbol is NotSpecified:
+                word_prefix_symbol = None
+        if word_prefix_symbol is not None:
+            # I'm not sure if this makes sense otherwise...
+            assert label_postfix_merge_symbol is None, (
+                f"{self}: word_prefix_symbol {word_prefix_symbol},"
+                f" label_postfix_merge_symbol {label_postfix_merge_symbol}"
+            )
 
         self.rnd = numpy.random.RandomState(0)
         self.bpe = SamplingBytePairEncoder(
@@ -342,6 +377,7 @@ class SamplingBytePairEncoding(Vocabulary):
             breadth_prob=breadth_prob,
             rnd=self.rnd,
             unknown_label=self.id_to_label(self.unknown_label_id) if self.unknown_label_id is not None else None,
+            opts=BpeOpts(label_postfix_merge_symbol=label_postfix_merge_symbol, word_prefix_symbol=word_prefix_symbol),
         )
 
     def set_random_seed(self, seed):

{returnn-1.20240620.64807 → returnn-1.20240620.201955}/returnn/util/bpe.py

@@ -4,6 +4,7 @@ Provide basic Byte-Pair-Encoding (BPE) utilities.
 
 from __future__ import annotations
 from typing import Optional, List, Dict, Callable
+from dataclasses import dataclass
 import re
 import numpy
 
@@ -227,70 +228,69 @@ class StandardBytePairEncoder:
         return output
 
 
+@dataclass
+class BpeOpts:
+    """
+    Options, should allow for both subword-nmt BPE and SentencePiece BPE/Unigram.
+    """
+
+    label_postfix_merge_symbol: Optional[str] = None  # eg. "@@"
+    word_prefix_symbol: Optional[str] = None  # eg. "▁"
+
+
 class PrefixTree:
     """
     Prefix tree / trie.
     This class represents both a single node and the tree.
     """
 
-    def __init__(self, prefix: str = "", root: Optional[PrefixTree] = None):
+    def __init__(self, *, prefix: str = "", opts: BpeOpts):
         """
-        :param prefix:
-        :param root:
+        :param prefix: if this is not the root, the prefix to get here
         """
         self.prefix = prefix
         self.arcs: Dict[str, PrefixTree] = {}  # single char (or BpePostMergeSymbol) -> sub tree
-        self.finished = False  # word finished here
+        self.finished = False  # label finished here
         self.bpe_finished = False  # partial word finished here with BpePostMergeSymbol at end
-        self.is_root = not root
-        self.root = root
+        self.opts = opts
 
-    def add(self, postfix: str, root: Optional[PrefixTree] = None) -> PrefixTree:
+    def add(self, postfix: str) -> PrefixTree:
         """
         :param postfix:
-        :param root:
         """
-        if not root:
-            if self.is_root:
-                root = self
-            else:
-                assert self.root
-                root = self.root
         self_ = self
+        opts = self.opts
+        label_postfix_merge_symbol = self.opts.label_postfix_merge_symbol
         while True:
-            if postfix == BpePostMergeSymbol:
+            if self_ is self and opts.word_prefix_symbol and postfix.startswith(opts.word_prefix_symbol):
+                arc = opts.word_prefix_symbol
+            elif postfix == label_postfix_merge_symbol:
                 arc = postfix
-                postfix_ = ""
             else:
                 arc = postfix[:1]
-                postfix_ = postfix[1:]
+            postfix_ = postfix[len(arc) :]
             if arc in self_.arcs:
                 child = self_.arcs[arc]
             else:
-                child = PrefixTree(root=root, prefix=self_.prefix + arc)
+                child = PrefixTree(prefix=self_.prefix + arc, opts=opts)
                 self_.arcs[arc] = child
-            if arc == BpePostMergeSymbol and not postfix_:
+            if arc == label_postfix_merge_symbol and not postfix_:
                 self_.bpe_finished = True
-            if postfix_:
-                self_ = child
-                postfix = postfix_
-            else:
+            if not postfix_:
                 child.finished = True
                 return child
+            self_ = child
+            postfix = postfix_
 
 
+@dataclass
 class Hyp:
     """
     Represents a hypothesis in the search.
     """
 
-    def __init__(self, bpe_sym_history, cur_node):
-        """
-        :param list[str] bpe_sym_history:
-        :param PrefixTree cur_node:
-        """
-        self.bpe_sym_history = bpe_sym_history
-        self.cur_node = cur_node
+    bpe_sym_history: List[str]
+    cur_node: PrefixTree
 
 
 class CharSyncSearch:
@@ -305,9 +305,15 @@ class CharSyncSearch:
         :param word_pos:
         """
         self.bpe = bpe
+        self.opts = bpe.opts
         self.word = word
         self.word_pos = word_pos
-        self.hyps: List[Hyp] = [Hyp(bpe_sym_history=[], cur_node=bpe)]
+        self.hyps: List[Hyp] = [
+            Hyp(
+                bpe_sym_history=[],
+                cur_node=bpe if self.opts.word_prefix_symbol is None else bpe.arcs[self.opts.word_prefix_symbol],
+            )
+        ]
         self.final_bpe_seqs: Optional[List[List[str]]] = None
 
     def _get_finished(self):
@@ -329,10 +335,17 @@ class CharSyncSearch:
                 if next_node:
                     new_hyps.append(
                         Hyp(
-                            bpe_sym_history=hyp.bpe_sym_history + [hyp.cur_node.prefix + BpePostMergeSymbol],
+                            bpe_sym_history=hyp.bpe_sym_history
+                            + [hyp.cur_node.prefix + self.opts.label_postfix_merge_symbol],
                             cur_node=next_node,
                         )
                     )
+            if self.opts.word_prefix_symbol is not None and hyp.cur_node.finished:
+                next_node = self.bpe.arcs.get(char)
+                if next_node:
+                    new_hyps.append(
+                        Hyp(bpe_sym_history=hyp.bpe_sym_history + [hyp.cur_node.prefix], cur_node=next_node)
+                    )
             next_node = hyp.cur_node.arcs.get(char)
             if next_node:
                 new_hyps.append(Hyp(bpe_sym_history=hyp.bpe_sym_history, cur_node=next_node))
@@ -349,20 +362,15 @@ class CharSyncSearch:
         return self.final_bpe_seqs
 
 
+@dataclass
 class HypInPos:
     """
     Represents a hypothesis in the search.
     """
 
-    def __init__(self, bpe_sym_history, cur_node, pos):
-        """
-        :param list[str] bpe_sym_history:
-        :param PrefixTree cur_node:
-        :param int pos:
-        """
-        self.bpe_sym_history = bpe_sym_history
-        self.cur_node = cur_node
-        self.pos = pos
+    bpe_sym_history: List[str]
+    cur_node: PrefixTree
+    pos: int
 
 
 class DepthFirstSearch:
@@ -377,11 +385,18 @@ class DepthFirstSearch:
         :param sampler:
         """
         self.bpe = bpe
+        self.opts = bpe.opts
         self.word = word
         self.sampler = sampler
         self.hyps: List[HypInPos] = []
         self.final_bpe_seq: Optional[List[str]] = None
-        self._add_hyp(HypInPos(bpe_sym_history=[], cur_node=bpe, pos=0))
+        self._add_hyp(
+            HypInPos(
+                bpe_sym_history=[],
+                cur_node=bpe if self.opts.word_prefix_symbol is None else bpe.arcs[self.opts.word_prefix_symbol],
+                pos=0,
+            )
+        )
 
     def _add_hyp(self, hyp: HypInPos):
         if hyp.pos >= len(self.word):
@@ -402,16 +417,26 @@ class DepthFirstSearch:
             if next_node:
                 new_hyps.append(
                     HypInPos(
-                        bpe_sym_history=hyp.bpe_sym_history + [hyp.cur_node.prefix + BpePostMergeSymbol],
+                        bpe_sym_history=hyp.bpe_sym_history
+                        + [hyp.cur_node.prefix + self.opts.label_postfix_merge_symbol],
                         cur_node=next_node,
                         pos=hyp.pos + 1,
                     )
                 )
+        if self.opts.word_prefix_symbol is not None and hyp.cur_node.finished:
+            next_node = self.bpe.arcs.get(char)
+            if next_node:
+                new_hyps.append(
+                    HypInPos(
+                        bpe_sym_history=hyp.bpe_sym_history + [hyp.cur_node.prefix], cur_node=next_node, pos=hyp.pos + 1
+                    )
+                )
         next_node = hyp.cur_node.arcs.get(char)
         if next_node:
             new_hyps.append(HypInPos(bpe_sym_history=hyp.bpe_sym_history, cur_node=next_node, pos=hyp.pos + 1))
 
         # Note that the order we check them will make this a depth-first or breadth-first search.
+        # We pop(-1) from the hyps list. So the last entry we add here will be the next one we expand.
         if self.sampler and self.sampler():
             new_hyps = list(reversed(new_hyps))
         for hyp in new_hyps:
@@ -431,13 +456,22 @@ class SamplingBytePairEncoder:
     Will randomly sample from any possible BPE split.
     """
 
-    def __init__(self, labels, breadth_prob, rnd, unknown_label=None):
-        """
-        :param list[str] labels: vocab
-        :param float breadth_prob: 1.0 will lead to breadth-first search, 0.0 to depth-first search.
-          other values are stochastic.
-        :param numpy.random.RandomState rnd:
-        :param str|None unknown_label:
+    def __init__(
+        self,
+        *,
+        labels: List[str],
+        breadth_prob: float,
+        rnd: numpy.random.RandomState,
+        unknown_label: Optional[str] = None,
+        opts: BpeOpts,
+    ):
+        """
+        :param labels: vocab
+        :param breadth_prob: 1.0 will lead to breadth-first search, 0.0 to depth-first search.
+            other values are stochastic.
+        :param rnd:
+        :param unknown_label:
+        :param opts:
         """
         self.labels = labels
         self.unknown_label = unknown_label
@@ -445,9 +479,10 @@ class SamplingBytePairEncoder:
             assert unknown_label in self.labels
         self.breadth_prob = breadth_prob
         self.rnd = rnd
+        self.opts = opts
 
         # build prefix tree
-        bpe = PrefixTree()
+        bpe = PrefixTree(opts=opts)
         for bpe_sym in labels:
             bpe.add(bpe_sym)
         self._bpe_prefix_tree = bpe
@@ -514,9 +549,10 @@ def _demo():
 
     vocab = Vocabulary(vocab_file=args.vocab, unknown_label=None)
     rnd = numpy.random.RandomState(args.seed)
+    opts = BpeOpts(label_postfix_merge_symbol=BpePostMergeSymbol)
 
     if args.input:
-        bpe_prefix_tree = PrefixTree()
+        bpe_prefix_tree = PrefixTree(opts=opts)
         for bpe_sym in vocab.labels:
             bpe_prefix_tree.add(bpe_sym)
 
@@ -533,7 +569,9 @@ def _demo():
                 print("%s: %s" % (word, " ".join(greedy)))
         return
 
-    bpe = SamplingBytePairEncoder(labels=vocab.labels, breadth_prob=args.breadth_prob, rnd=rnd, unknown_label=args.unk)
+    bpe = SamplingBytePairEncoder(
+        labels=vocab.labels, breadth_prob=args.breadth_prob, rnd=rnd, unknown_label=args.unk, opts=opts
+    )
     print("Reading from stdin:")
     while True:
         try:

{returnn-1.20240620.64807 → returnn-1.20240620.201955/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240620.64807 → returnn-1.20240620.201955}/tests/test_Util.py

@@ -638,9 +638,9 @@ def test_expand_env_vars():
 
 
 def test_bpe_PrefixTree():
-    from returnn.util.bpe import PrefixTree, BpePostMergeSymbol
+    from returnn.util.bpe import PrefixTree, BpeOpts, BpePostMergeSymbol
 
-    tree = PrefixTree()
+    tree = PrefixTree(opts=BpeOpts(label_postfix_merge_symbol=BpePostMergeSymbol))
     tree.add("hello")
     tree.add("helo" + BpePostMergeSymbol)
     assert not tree.finished and not tree.bpe_finished
@@ -663,6 +663,136 @@ def test_bpe_PrefixTree():
     assert set(node_o_post.arcs.keys()) == set()
 
 
+def test_bpe_PrefixTree_word_prefix():
+    from returnn.util.bpe import PrefixTree, BpeOpts
+
+    tree = PrefixTree(opts=BpeOpts(word_prefix_symbol="▁"))
+    tree.add("▁hello")
+    tree.add("▁hel")
+    tree.add("lo")
+    tree.add("▁hi")
+    assert not tree.finished and not tree.bpe_finished
+    assert set(tree.arcs.keys()) == {"▁", "l"}
+    node = tree.arcs["▁"]
+    assert not node.finished and not node.bpe_finished
+    assert set(node.arcs.keys()) == {"h"}
+    node = node.arcs["h"]
+    assert not node.finished and not node.bpe_finished
+    assert set(node.arcs.keys()) == {"e", "i"}
+    node_ = node.arcs["i"]
+    assert node_.finished and not node_.bpe_finished and not node_.arcs
+    node = node.arcs["e"]
+    assert not node.finished and not node.bpe_finished
+    assert set(node.arcs.keys()) == {"l"}
+    node = node.arcs["l"]
+    assert node.finished and not node.bpe_finished
+    assert set(node.arcs.keys()) == {"l"}
+    node = node.arcs["l"]
+    assert not node.finished and not node.bpe_finished
+    assert set(node.arcs.keys()) == {"o"}
+    node = node.arcs["o"]
+    assert node.finished and not node.bpe_finished and not node.arcs
+
+
+def test_bpe_DepthFirstSearch():
+    import itertools
+    from returnn.util.bpe import PrefixTree, BpeOpts, DepthFirstSearch
+
+    tree = PrefixTree(opts=BpeOpts(label_postfix_merge_symbol="@@"))
+    tree.add("llo")
+    tree.add("helo@@")
+    tree.add("he@@")
+
+    dfs = DepthFirstSearch(tree, "hello")
+    assert_equal(dfs.search(), ["he@@", "llo"])
+    dfs = DepthFirstSearch(tree, "helo")
+    assert_equal(dfs.search(), None)
+    dfs = DepthFirstSearch(tree, "x")
+    assert_equal(dfs.search(), None)
+    dfs = DepthFirstSearch(tree, "llo")
+    assert_equal(dfs.search(), ["llo"])
+
+    tree.add("hello")
+    dfs = DepthFirstSearch(tree, "hello")
+    assert_equal(dfs.search(), ["hello"])
+    dfs = DepthFirstSearch(tree, "hello", sampler=lambda: True)
+    assert_equal(dfs.search(), ["he@@", "llo"])
+
+    tree.add("hel@@")
+    tree.add("lo")
+    dfs = DepthFirstSearch(tree, "hello")
+    assert_equal(dfs.search(), ["hello"])
+    dfs = DepthFirstSearch(tree, "hello", sampler=lambda: True)
+    assert_equal(dfs.search(), ["he@@", "llo"])
+    dfs = DepthFirstSearch(tree, "hello", sampler=lambda _it=itertools.count(): next(_it) in {3})
+    assert_equal(dfs.search(), ["hel@@", "lo"])
+
+
+def test_bpe_DepthFirstSearch_word_prefix():
+    import itertools
+    from returnn.util.bpe import PrefixTree, BpeOpts, DepthFirstSearch
+
+    tree = PrefixTree(opts=BpeOpts(word_prefix_symbol="▁"))
+    tree.add("▁hello")
+    tree.add("▁hel")
+    tree.add("lo")
+
+    search = DepthFirstSearch(tree, "hello")
+    assert_equal(search.search(), ["▁hello"])
+    search = DepthFirstSearch(tree, "hello", sampler=lambda: True)
+    assert_equal(search.search(), ["▁hel", "lo"])
+
+    tree.add("▁he")
+    tree.add("llo")
+    search = DepthFirstSearch(tree, "hello", sampler=lambda: True)
+    assert_equal(search.search(), ["▁he", "llo"])
+    search = DepthFirstSearch(tree, "hello", sampler=lambda _it=itertools.count(): next(_it) in {3})
+    assert_equal(search.search(), ["▁hel", "lo"])
+
+
+def test_bpe_CharSyncSearch():
+    from returnn.util.bpe import PrefixTree, BpeOpts, CharSyncSearch
+
+    tree = PrefixTree(opts=BpeOpts(label_postfix_merge_symbol="@@"))
+    tree.add("llo")
+    tree.add("helo@@")
+    tree.add("he@@")
+
+    search = CharSyncSearch(tree, "hello")
+    assert_equal(search.search(), [["he@@", "llo"]])
+    search = CharSyncSearch(tree, "helo")
+    assert_equal(search.search(), [])
+    search = CharSyncSearch(tree, "x")
+    assert_equal(search.search(), [])
+    search = CharSyncSearch(tree, "llo")
+    assert_equal(search.search(), [["llo"]])
+
+    tree.add("hello")
+    search = CharSyncSearch(tree, "hello")
+    assert_equal(search.search(), [["he@@", "llo"], ["hello"]])
+
+
+def test_bpe_CharSyncSearch_word_prefix():
+    from returnn.util.bpe import PrefixTree, BpeOpts, CharSyncSearch
+
+    tree = PrefixTree(opts=BpeOpts(word_prefix_symbol="▁"))
+    tree.add("▁hello")
+    tree.add("▁hel")
+    tree.add("lo")
+    tree.add("▁hi")
+
+    search = CharSyncSearch(tree, "hello")
+    assert_equal(search.search(), [["▁hel", "lo"], ["▁hello"]])
+    search = CharSyncSearch(tree, "helo")
+    assert_equal(search.search(), [])
+    search = CharSyncSearch(tree, "x")
+    assert_equal(search.search(), [])
+    search = CharSyncSearch(tree, "lo")
+    assert_equal(search.search(), [])
+    search = CharSyncSearch(tree, "hi")
+    assert_equal(search.search(), [["▁hi"]])
+
+
 def test_file_cache():
     from returnn.util.file_cache import FileCache, CachedFile
 

{returnn-1.20240620.64807 → returnn-1.20240620.201955}/tools/bpe-create-lexicon.py

@@ -96,7 +96,7 @@ def main():
     print("Words: num %i, first %r" % (len(words), words[0]))
 
     print("Build BPE prefix tree...")
-    bpe = bpe_utils.PrefixTree()
+    bpe = bpe_utils.PrefixTree(opts=bpe_utils.BpeOpts(label_postfix_merge_symbol=bpe_utils.BpePostMergeSymbol))
     for bpe_sym in bpe_syms:
         bpe.add(bpe_sym)
 

returnn-1.20240620.64807/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240620.064807'
-long_version = '1.20240620.064807+git.ab33b05'