Repository: OpenNMT/OpenNMT-tf

Issue:
exported model does not include vocabulary assets
When using Version `2.1` and calling `onmt-main ... export` the exported model's `assets` directory is empty, rather than including BPE vocab files as in previous versions. Is this intended behaviour? I could submit a patch if it isn't.


Discussion:
Yes, it is the expected behavior. Because of this TensorFlow issue https://github.com/tensorflow/tensorflow/issues/32770, the vocabulary is currently embedded in the graph itself. I plan to revise this once the TensorFlow fix is released.

Resolving patch:
diff --git a/opennmt/inputters/text_inputter.py b/opennmt/inputters/text_inputter.py
--- a/opennmt/inputters/text_inputter.py
+++ b/opennmt/inputters/text_inputter.py
@@ -207,18 +207,35 @@ def _get_field(config, key, prefix=None, default=None, required=False):
     raise ValueError("Missing field '%s' in the data configuration" % key)
   return value
 
-def _create_tokens_to_ids_table(tokens, ids, num_oov_buckets):
-  # TODO: consider reverting back to TextFileInitializer when this change is released:
-  # https://github.com/tensorflow/tensorflow/pull/32773
-  initializer = tf.lookup.KeyValueTensorInitializer(tokens, ids)
+def _create_vocabulary_tables(vocabulary_file, num_oov_buckets, as_asset=True):
+  vocabulary = Vocab.from_file(vocabulary_file)
+  vocabulary_size = len(vocabulary)
+  if as_asset:
+    tokens_to_ids_initializer = tf.lookup.TextFileInitializer(
+        vocabulary_file,
+        tf.string,
+        tf.lookup.TextFileIndex.WHOLE_LINE,
+        tf.int64,
+        tf.lookup.TextFileIndex.LINE_NUMBER,
+        vocab_size=vocabulary_size)
+    ids_to_tokens_initializer = tf.lookup.TextFileInitializer(
+        vocabulary_file,
+        tf.int64,
+        tf.lookup.TextFileIndex.LINE_NUMBER,
+        tf.string,
+        tf.lookup.TextFileIndex.WHOLE_LINE,
+        vocab_size=vocabulary_size)
+  else:
+    tokens = tf.constant(vocabulary.words, dtype=tf.string)
+    ids = tf.constant(list(range(vocabulary_size)), dtype=tf.int64)
+    tokens_to_ids_initializer = tf.lookup.KeyValueTensorInitializer(tokens, ids)
+    ids_to_tokens_initializer = tf.lookup.KeyValueTensorInitializer(ids, tokens)
   if num_oov_buckets > 0:
-    return tf.lookup.StaticVocabularyTable(initializer, num_oov_buckets)
+    tokens_to_ids = tf.lookup.StaticVocabularyTable(tokens_to_ids_initializer, num_oov_buckets)
   else:
-    return tf.lookup.StaticHashTable(initializer, 0)
-
-def _create_ids_to_tokens_table(ids, tokens):
-  initializer = tf.lookup.KeyValueTensorInitializer(ids, tokens)
-  return tf.lookup.StaticHashTable(initializer, constants.UNKNOWN_TOKEN)
+    tokens_to_ids = tf.lookup.StaticHashTable(tokens_to_ids_initializer, 0)
+  ids_to_tokens = tf.lookup.StaticHashTable(ids_to_tokens_initializer, constants.UNKNOWN_TOKEN)
+  return vocabulary_size + num_oov_buckets, tokens_to_ids, ids_to_tokens
 
 
 class TextInputter(Inputter):
@@ -234,12 +251,10 @@ def __init__(self, num_oov_buckets=1, **kwargs):
   def initialize(self, data_config, asset_prefix=""):
     self.vocabulary_file = _get_field(
         data_config, "vocabulary", prefix=asset_prefix, required=True)
-    vocabulary = Vocab.from_file(self.vocabulary_file)
-    self.vocabulary_size = len(vocabulary) + self.num_oov_buckets
-    tokens = tf.constant(vocabulary.words, dtype=tf.string)
-    ids = tf.constant(list(range(len(vocabulary))), dtype=tf.int64)
-    self.tokens_to_ids = _create_tokens_to_ids_table(tokens, ids, self.num_oov_buckets)
-    self.ids_to_tokens = _create_ids_to_tokens_table(ids, tokens)
+    self.vocabulary_size, self.tokens_to_ids, self.ids_to_tokens = _create_vocabulary_tables(
+        self.vocabulary_file,
+        self.num_oov_buckets,
+        as_asset=data_config.get("export_vocabulary_assets", True))
     tokenizer_config = _get_field(data_config, "tokenization", prefix=asset_prefix)
     self.tokenizer = tokenizers.make_tokenizer(tokenizer_config)
 


Test patch:
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -233,9 +233,11 @@ def testScore(self):
       lines = f.readlines()
     self.assertEqual(len(lines), 5)
 
-  def testExport(self):
+  @parameterized.expand([[True], [False]])
+  def testExport(self, export_vocabulary_assets):
     config = {
         "data": {
+            "export_vocabulary_assets": export_vocabulary_assets,
             "source_tokenization": {
                 "mode": "char"
             }
@@ -245,10 +247,23 @@ def testExport(self):
     runner = self._getTransliterationRunner(config)
     runner.export(export_dir)
     self.assertTrue(tf.saved_model.contains_saved_model(export_dir))
+
+    # Check assets directories.
+    assets = os.listdir(os.path.join(export_dir, "assets"))
+    if export_vocabulary_assets:
+      self.assertLen(assets, 2)
+    else:
+      self.assertLen(assets, 0)
     extra_assets_dir = os.path.join(export_dir, "assets.extra")
     self.assertTrue(os.path.isdir(extra_assets_dir))
     self.assertLen(os.listdir(extra_assets_dir), 1)
-    imported = tf.saved_model.load(export_dir)
+
+    # Export directory could be relocated and does not reference the original vocabulary files.
+    shutil.rmtree(runner.model_dir)
+    export_dir_2 = os.path.join(self.get_temp_dir(), "export_2")
+    os.rename(export_dir, export_dir_2)
+    self.assertTrue(tf.saved_model.contains_saved_model(export_dir_2))
+    imported = tf.saved_model.load(export_dir_2)
     translate_fn = imported.signatures["serving_default"]
     outputs = translate_fn(
         tokens=tf.constant([["آ" ,"ت" ,"ز" ,"م" ,"و" ,"ن"]]),