returnn 1.20240619.172413__tar.gz → 1.20240619.184216__tar.gz

{returnn-1.20240619.172413/returnn.egg-info → returnn-1.20240619.184216}/PKG-INFO

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

returnn-1.20240619.184216/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240619.184216'
+long_version = '1.20240619.184216+git.bc0e997'

{returnn-1.20240619.172413 → returnn-1.20240619.184216}/returnn/util/bpe.py

@@ -2,12 +2,14 @@
 Provide basic Byte-Pair-Encoding (BPE) utilities.
 """
 
-import re
+from __future__ import annotations
+from typing import Dict
 import typing
+import re
 import numpy
 
 
-BpeMergeSymbol = "@@"
+BpePostMergeSymbol = "@@"
 
 
 class StandardBytePairEncoder:
@@ -33,7 +35,7 @@ class StandardBytePairEncoder:
         self.labels = labels
         self._load(bpe_codes_file)
         self._bpe_encode_cache = {}
-        self._bpe_separator = BpeMergeSymbol
+        self._bpe_separator = BpePostMergeSymbol
 
     _file_cache = {}  # filename -> bpe_file_version, bpe_codes, bpe_codes_reverse
 
@@ -238,9 +240,9 @@ class PrefixTree:
         :param PrefixTree|None root:
         """
         self.prefix = prefix
-        self.arcs = {}  # type: typing.Dict[str,PrefixTree]
-        self.finished = False
-        self.bpe_finished = False
+        self.arcs: Dict[str, PrefixTree] = {}  # single char (or BpePostMergeSymbol) -> sub tree
+        self.finished = False  # word finished here
+        self.bpe_finished = False  # partial word finished here with BpePostMergeSymbol at end
         self.is_root = not root
         self.root = root
 
@@ -256,24 +258,27 @@ class PrefixTree:
             else:
                 assert self.root
                 root = self.root
-        if postfix == BpeMergeSymbol:
-            arc = postfix
-            postfix_ = ""
-        else:
-            arc = postfix[:1]
-            postfix_ = postfix[1:]
-        if arc in self.arcs:
-            child = self.arcs[arc]
-        else:
-            child = PrefixTree(root=root, prefix=self.prefix + arc)
-            self.arcs[arc] = child
-        if arc == BpeMergeSymbol and not postfix_:
-            self.bpe_finished = True
-        if postfix_:
-            return child.add(postfix_, root=root)
-        else:
-            child.finished = True
-            return child
+        self_ = self
+        while True:
+            if postfix == BpePostMergeSymbol:
+                arc = postfix
+                postfix_ = ""
+            else:
+                arc = postfix[:1]
+                postfix_ = postfix[1:]
+            if arc in self_.arcs:
+                child = self_.arcs[arc]
+            else:
+                child = PrefixTree(root=root, prefix=self_.prefix + arc)
+                self_.arcs[arc] = child
+            if arc == BpePostMergeSymbol and not postfix_:
+                self_.bpe_finished = True
+            if postfix_:
+                self_ = child
+                postfix = postfix_
+            else:
+                child.finished = True
+                return child
 
 
 class Hyp:
@@ -321,11 +326,12 @@ class CharSyncSearch:
         new_hyps = []  # type: typing.List[Hyp]
         for hyp in self.hyps:
             if hyp.cur_node.bpe_finished:
+                # Start again from root node.
                 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 + BpeMergeSymbol],
+                            bpe_sym_history=hyp.bpe_sym_history + [hyp.cur_node.prefix + BpePostMergeSymbol],
                             cur_node=next_node,
                         )
                     )
@@ -397,11 +403,12 @@ class DepthFirstSearch:
         char = self.word[hyp.pos]
         new_hyps = []  # type: typing.List[HypInPos]
         if hyp.cur_node.bpe_finished:
+            # Start again from root node.
             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 + BpeMergeSymbol],
+                        bpe_sym_history=hyp.bpe_sym_history + [hyp.cur_node.prefix + BpePostMergeSymbol],
                         cur_node=next_node,
                         pos=hyp.pos + 1,
                     )

{returnn-1.20240619.172413 → returnn-1.20240619.184216/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240619.172413 → returnn-1.20240619.184216}/tests/test_Util.py

@@ -637,6 +637,32 @@ def test_expand_env_vars():
     assert_equal(expand_env_vars("$TMPA/$TMPB/returnn/file_cache"), "/testA/testB/returnn/file_cache")
 
 
+def test_bpe_PrefixTree():
+    from returnn.util.bpe import PrefixTree, BpePostMergeSymbol
+
+    tree = PrefixTree()
+    tree.add("hello")
+    tree.add("helo" + BpePostMergeSymbol)
+    assert not tree.finished and not tree.bpe_finished
+    assert set(tree.arcs.keys()) == {"h"}
+    node = tree.arcs["h"]
+    assert not node.finished and not node.bpe_finished
+    assert set(node.arcs.keys()) == {"e"}
+    node = node.arcs["e"]
+    assert not 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()) == {"l", "o"}
+    node_o = node.arcs["o"]
+    assert not node_o.finished and node_o.bpe_finished
+    # The following BpePostMergeSymbol arc and its node is somewhat arbitrary, probably does not matter...
+    assert set(node_o.arcs.keys()) == {BpePostMergeSymbol}
+    node_o_post = node_o.arcs[BpePostMergeSymbol]
+    assert node_o_post.finished and not node_o_post.bpe_finished
+    assert set(node_o_post.arcs.keys()) == set()
+
+
 def test_file_cache():
     from returnn.util.file_cache import FileCache, CachedFile
 

{returnn-1.20240619.172413 → returnn-1.20240619.184216}/tools/bpe-create-lexicon.py

@@ -140,7 +140,7 @@ def main():
         visit_word(word)
 
     for bpe_sym in bpe_syms:
-        if bpe_sym.endswith(bpe_utils.BpeMergeSymbol):
+        if bpe_sym.endswith(bpe_utils.BpePostMergeSymbol):
             continue
         if bpe_sym not in words:
             continue

returnn-1.20240619.172413/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240619.172413'
-long_version = '1.20240619.172413+git.d5048fd'