Repository: awslabs/sockeye

Issue:
bug when re-loading 10 or more source vocabs
I believe there is a bug in `vocab.load_source_vocabs` that is apparent only when 9 or more source-factors are in use (10+ total source vocabularies). The prepared source vocabularies for each of the source-factors are ordered by integer labels (`vocab.src.0.json`, `vocab.src.1.json`, `vocab.src.2.json` etc..) However, when these source vocabs are loaded from a sockeye-prepared data directory, they are sorted in alphabetic/string order, which means that `vocab.src.10.json` will come before `vocab.src.2.json` and so forth. This means that the vocabularies originally prepared by sockeye.prepare-data will not match the re-loaded vocabularies (see error below).

```
def load_source_vocabs(folder: str) -> List[Vocab]:
    """
    Loads source vocabularies from folder. The first element in the list is the primary source vocabulary.
    Other elements correspond to optional additional source factor vocabularies found in folder.
    :param folder: Source folder.
    :return: List of vocabularies.
    """
    return [vocab_from_json(os.path.join(folder, fname)) for fname in
            sorted([f for f in os.listdir(folder) if f.startswith(C.VOCAB_SRC_PREFIX)])]
```

Specifically, the issue arises when I try to resume training on a model where I'm using the `--prepared-data` flag, giving the following error:

```
Traceback (most recent call last):
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/site-packages/sockeye/train.py", line 966, in <module>
    main()
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/site-packages/sockeye/train.py", line 832, in main
    train(args)
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/site-packages/sockeye/train.py", line 886, in train
    output_folder=output_folder)
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/site-packages/sockeye/train.py", line 271, in create_data_iters_and_vocabs
    "Prepared data and resumed model source vocab %d do not match." % i)
  File "/home/rrudinger/.conda/envs/sockeye10/lib/python3.6/site-packages/sockeye/utils.py", line 121, in check_condition
    raise SockeyeError(error_message)
sockeye.utils.SockeyeError: Prepared data and resumed model source vocab 2 do not match.
```

Edit: Assuming that all sockeye-created source vocabs are of the form `*.N.json` where N is an integer, then I think the following would fix the issue:

```
def load_source_vocabs(folder: str) -> List[Vocab]:
    """
    Loads source vocabularies from folder. The first element in the list is the primary source vocabulary.
    Other elements correspond to optional additional source factor vocabularies found in folder.
    :param folder: Source folder.
    :return: List of vocabularies.
    """
    return [vocab_from_json(os.path.join(folder, fname)) for fname in
            sorted([f for f in os.listdir(folder) if f.startswith(C.VOCAB_SRC_PREFIX)], key=lambda x: int(x.split('.')[-2]))]

```


Discussion:
Great catch, thank you! Do you want to open a pull request with your proposed change?

Resolving patch:
diff --git a/sockeye/__init__.py b/sockeye/__init__.py
--- a/sockeye/__init__.py
+++ b/sockeye/__init__.py
@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '1.18.101'
+__version__ = '1.18.102'
diff --git a/sockeye/vocab.py b/sockeye/vocab.py
--- a/sockeye/vocab.py
+++ b/sockeye/vocab.py
@@ -178,6 +178,11 @@ def save_target_vocab(target_vocab: Vocab, folder: str):
     vocab_to_json(target_vocab, os.path.join(folder, C.VOCAB_TRG_NAME % 0))
 
 
+def _get_sorted_source_vocab_fnames(folder) -> List[str]:
+    _key = lambda x: int(x.split('.', 3)[-2])
+    return sorted([f for f in os.listdir(folder) if f.startswith(C.VOCAB_SRC_PREFIX)], key=_key)
+
+
 def load_source_vocabs(folder: str) -> List[Vocab]:
     """
     Loads source vocabularies from folder. The first element in the list is the primary source vocabulary.
@@ -186,8 +191,7 @@ def load_source_vocabs(folder: str) -> List[Vocab]:
     :param folder: Source folder.
     :return: List of vocabularies.
     """
-    return [vocab_from_json(os.path.join(folder, fname)) for fname in
-            sorted([f for f in os.listdir(folder) if f.startswith(C.VOCAB_SRC_PREFIX)])]
+    return [vocab_from_json(os.path.join(folder, fname)) for fname in _get_sorted_source_vocab_fnames(folder)]
 
 
 def load_target_vocab(folder: str) -> Vocab:


Test patch:
diff --git a/test/unit/test_vocab.py b/test/unit/test_vocab.py
--- a/test/unit/test_vocab.py
+++ b/test/unit/test_vocab.py
@@ -12,9 +12,10 @@
 # permissions and limitations under the License.
 
 import pytest
+from unittest import mock
 
 import sockeye.constants as C
-from sockeye.vocab import build_vocab, get_ordered_tokens_from_vocab, is_valid_vocab
+from sockeye.vocab import build_vocab, get_ordered_tokens_from_vocab, is_valid_vocab, _get_sorted_source_vocab_fnames
 
 test_vocab = [
         # Example 1
@@ -100,3 +101,11 @@ def test_get_ordered_tokens_from_vocab(vocab, expected_output):
 )
 def test_verify_valid_vocab(vocab, expected_result):
     assert is_valid_vocab(vocab) == expected_result
+
+
+def test_get_sorted_source_vocab_fnames():
+    expected_fnames = [C.VOCAB_SRC_NAME % i for i in [1, 2, 10]]
+    with mock.patch('os.listdir') as mocked_listdir:
+        mocked_listdir.return_value = [C.VOCAB_SRC_NAME % i for i in [2, 1, 10]]
+        fnames = _get_sorted_source_vocab_fnames(None)
+        assert fnames == expected_fnames