Repository: OpenNMT/OpenNMT-tf
Pull request #888 (2021-10-05 13:56:40)

Add intermediate methods format_{prediction,score}


Patch:
diff --git a/config/models/multi_source_transformer.py b/config/models/multi_source_transformer.py
index 89af985a..59cf68bc 100644
--- a/config/models/multi_source_transformer.py
+++ b/config/models/multi_source_transformer.py
@@ -17,8 +17,6 @@ data:
 
 import opennmt
 
-from opennmt.utils import misc
-
 
 class DualSourceTransformer(opennmt.models.Transformer):
     def __init__(self):
@@ -43,7 +41,7 @@ class DualSourceTransformer(opennmt.models.Transformer):
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
         max_length = config["train"]["maximum_features_length"]
-        return misc.merge_dict(
+        return opennmt.merge_config(
             config, {"train": {"maximum_features_length": [max_length, max_length]}}
         )
 
diff --git a/opennmt/__init__.py b/opennmt/__init__.py
index e423e055..355bf5c7 100644
--- a/opennmt/__init__.py
+++ b/opennmt/__init__.py
@@ -4,7 +4,7 @@ from opennmt.version import __version__, _check_tf_version
 
 _check_tf_version()
 
-from opennmt.config import convert_to_v2_config, load_config, load_model
+from opennmt.config import convert_to_v2_config, load_config, load_model, merge_config
 from opennmt.constants import (
     END_OF_SENTENCE_ID,
     END_OF_SENTENCE_TOKEN,
diff --git a/opennmt/config.py b/opennmt/config.py
index e8705d49..0e217e50 100644
--- a/opennmt/config.py
+++ b/opennmt/config.py
@@ -152,11 +152,17 @@ def load_config(config_paths, config=None):
         with tf.io.gfile.GFile(config_path) as config_file:
             subconfig = yaml.safe_load(config_file.read())
             # Add or update section in main configuration.
-            merge_dict(config, subconfig)
+            merge_config(config, subconfig)
 
     return config
 
 
+def merge_config(a, b):
+    """Merges the configuration :obj:`b` into the configuration :obj:`a`."""
+    override_keys = {"optimizer_params", "decay_params"}
+    return merge_dict(a, b, override_keys=override_keys)
+
+
 def _map_config_values(config, fn):
     """Applies the function ``fn`` on each value of the configuration."""
     if isinstance(config, dict):
diff --git a/opennmt/models/catalog.py b/opennmt/models/catalog.py
index ef32074a..cd611494 100644
--- a/opennmt/models/catalog.py
+++ b/opennmt/models/catalog.py
@@ -3,6 +3,7 @@
 import tensorflow as tf
 import tensorflow_addons as tfa
 
+from opennmt import config as config_util
 from opennmt import decoders, encoders, inputters, layers
 from opennmt.models import (
     language_model,
@@ -75,7 +76,7 @@ class ListenAttendSpell(sequence_to_sequence.SequenceToSequence):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return misc.merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "params": {
@@ -99,7 +100,7 @@ class _RNNBase(sequence_to_sequence.SequenceToSequence):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return misc.merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "params": {
@@ -274,7 +275,7 @@ class LstmCnnCrfTagger(sequence_tagger.SequenceTagger):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return misc.merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "params": {
@@ -397,7 +398,7 @@ class GPT2Small(language_model.LanguageModel):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return misc.merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "params": {
diff --git a/opennmt/models/language_model.py b/opennmt/models/language_model.py
index c688ddf2..00238e7c 100644
--- a/opennmt/models/language_model.py
+++ b/opennmt/models/language_model.py
@@ -2,6 +2,7 @@
 
 import tensorflow as tf
 
+from opennmt import config as config_util
 from opennmt import inputters
 from opennmt.models import model
 from opennmt.utils import decoding, losses, misc
@@ -30,7 +31,7 @@ class LanguageModel(model.SequenceGenerator):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return misc.merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "infer": {
@@ -137,12 +138,11 @@ class LanguageModel(model.SequenceGenerator):
             training=training,
         )
 
-    def print_prediction(self, prediction, params=None, stream=None):
+    def format_prediction(self, prediction, params=None):
         target_length = prediction["length"]
         tokens = prediction["tokens"][:target_length]
         sentence = self.examples_inputter.tokenizer.detokenize(tokens)
-        sentence = misc.format_translation_output(sentence)
-        misc.print_as_bytes(sentence, stream=stream)
+        return misc.format_translation_output(sentence)
 
 
 class LanguageModelInputter(inputters.WordEmbedder, inputters.ExampleInputterAdapter):
diff --git a/opennmt/models/model.py b/opennmt/models/model.py
index fed9c293..55c695ee 100644
--- a/opennmt/models/model.py
+++ b/opennmt/models/model.py
@@ -500,27 +500,57 @@ class Model(tf.keras.layers.Layer):
         if not self.unsupervised:
             self.labels_inputter.visualize(self, log_dir)
 
+    def format_prediction(self, prediction, params=None):
+        """Formats the model prediction for file saving.
+
+        Args:
+          prediction: The model prediction (same structure as the second output of
+            :meth:`opennmt.models.Model.call`).
+          params: (optional) Dictionary of formatting parameters.
+
+        Returns:
+          A string or list of strings.
+        """
+        return str(prediction)
+
+    def format_score(self, score, params=None, stream=None):
+        """Formats the score result for file saving.
+
+        Args:
+          score: The score result (same structure as the output of
+            :meth:`opennmt.models.Model.score`).
+          params: (optional) Dictionary of formatting parameters.
+        """
+        return str(score)
+
     def print_prediction(self, prediction, params=None, stream=None):
         """Prints the model prediction.
 
         Args:
-          prediction: The model prediction.
+          prediction: The model prediction (same structure as the second output of
+            :meth:`opennmt.models.Model.call`).
           params: (optional) Dictionary of formatting parameters.
           stream: (optional) The stream to print to.
         """
-        _ = params
-        print(prediction, file=stream)
+        _write_lines(self.format_prediction(prediction, params=params), stream)
 
     def print_score(self, score, params=None, stream=None):
         """Prints the score result.
 
         Args:
-          score: The score result (output of :meth:`opennmt.models.Model.score`).
+          score: The score result (same structure as the output of
+            :meth:`opennmt.models.Model.score`).
           params: (optional) Dictionary of formatting parameters.
           stream: (optional) The stream to print to.
         """
-        _ = params
-        print(score, file=stream)
+        _write_lines(self.format_score(score, params=params), stream)
+
+
+def _write_lines(lines, stream):
+    if not isinstance(lines, list):
+        lines = [lines]
+    for line in lines:
+        misc.print_as_bytes(line, stream=stream)
 
 
 class SequenceGenerator(Model):
@@ -553,7 +583,7 @@ class SequenceGenerator(Model):
                 results[key_to_forward] = value
         return results
 
-    def print_score(self, score, params=None, stream=None):
+    def format_score(self, score, params=None):
         if params is None:
             params = {}
         length = score["length"]
@@ -566,11 +596,10 @@ class SequenceGenerator(Model):
         if "attention" in score:
             attention = score["attention"][:length]
         alignment_type = params.get("with_alignments")
-        sentence = misc.format_translation_output(
+        return misc.format_translation_output(
             sentence,
             score=score["score"],
             token_level_scores=token_level_scores,
             attention=attention,
             alignment_type=alignment_type,
         )
-        misc.print_as_bytes(sentence, stream=stream)
diff --git a/opennmt/models/sequence_classifier.py b/opennmt/models/sequence_classifier.py
index b5af0e60..e22670fb 100644
--- a/opennmt/models/sequence_classifier.py
+++ b/opennmt/models/sequence_classifier.py
@@ -4,7 +4,6 @@ import tensorflow as tf
 
 from opennmt import inputters
 from opennmt.models.model import Model
-from opennmt.utils import misc
 from opennmt.utils.losses import cross_entropy_loss
 
 
@@ -78,8 +77,8 @@ class SequenceClassifier(Model):
             labels["classes_id"], predictions["classes_id"]
         )
 
-    def print_prediction(self, prediction, params=None, stream=None):
-        misc.print_as_bytes(prediction["classes"], stream=stream)
+    def format_prediction(self, prediction, params=None):
+        return prediction["classes"].decode("utf-8")
 
 
 class ClassInputter(inputters.TextInputter):
diff --git a/opennmt/models/sequence_tagger.py b/opennmt/models/sequence_tagger.py
index 822fbb7c..2b2b7780 100644
--- a/opennmt/models/sequence_tagger.py
+++ b/opennmt/models/sequence_tagger.py
@@ -6,7 +6,6 @@ import tensorflow_addons as tfa
 
 from opennmt import inputters
 from opennmt.models.model import Model
-from opennmt.utils import misc
 from opennmt.utils.losses import cross_entropy_sequence_loss
 
 
@@ -121,10 +120,9 @@ class SequenceTagger(Model):
 
             metrics["f1"].update_state(gold_flags, predicted_flags)
 
-    def print_prediction(self, prediction, params=None, stream=None):
+    def format_prediction(self, prediction, params=None):
         tags = prediction["tags"][: prediction["length"]]
-        sent = b" ".join(tags)
-        misc.print_as_bytes(sent, stream=stream)
+        return " ".join(tag.decode("utf-8") for tag in tags)
 
 
 class TagsInputter(inputters.TextInputter):
diff --git a/opennmt/models/sequence_to_sequence.py b/opennmt/models/sequence_to_sequence.py
index 180e27cf..09f4d7e6 100644
--- a/opennmt/models/sequence_to_sequence.py
+++ b/opennmt/models/sequence_to_sequence.py
@@ -3,6 +3,7 @@
 import tensorflow as tf
 import tensorflow_addons as tfa
 
+from opennmt import config as config_util
 from opennmt import constants, inputters
 from opennmt.data import noise, text, vocab
 from opennmt.decoders import decoder as decoder_util
@@ -99,7 +100,7 @@ class SequenceToSequence(model.SequenceGenerator):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return misc.merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "params": {
@@ -458,7 +459,7 @@ class SequenceToSequence(model.SequenceGenerator):
                     )
         return loss, loss_normalizer, loss_token_normalizer
 
-    def print_prediction(self, prediction, params=None, stream=None):
+    def format_prediction(self, prediction, params=None):
         if params is None:
             params = {}
         with_scores = params.get("with_scores")
@@ -468,6 +469,7 @@ class SequenceToSequence(model.SequenceGenerator):
                 "with_alignments is set but the model did not return alignment information"
             )
         num_hypotheses = params.get("n_best", len(prediction["log_probs"]))
+        outputs = []
         for i in range(num_hypotheses):
             if "tokens" in prediction:
                 target_length = prediction["length"][i]
@@ -487,7 +489,8 @@ class SequenceToSequence(model.SequenceGenerator):
                 attention=attention,
                 alignment_type=alignment_type,
             )
-            misc.print_as_bytes(sentence, stream=stream)
+            outputs.append(sentence)
+        return outputs
 
     def transfer_weights(
         self, new_model, new_optimizer=None, optimizer=None, ignore_weights=None
diff --git a/opennmt/models/transformer.py b/opennmt/models/transformer.py
index 66b20a05..fbd70146 100644
--- a/opennmt/models/transformer.py
+++ b/opennmt/models/transformer.py
@@ -2,6 +2,7 @@
 
 import tensorflow as tf
 
+from opennmt import config as config_util
 from opennmt import inputters
 from opennmt.decoders.self_attention_decoder import SelfAttentionDecoder
 from opennmt.encoders.encoder import ParallelEncoder
@@ -12,7 +13,6 @@ from opennmt.models.sequence_to_sequence import (
     EmbeddingsSharingLevel,
     SequenceToSequence,
 )
-from opennmt.utils.misc import merge_dict
 
 
 class Transformer(SequenceToSequence):
@@ -173,7 +173,7 @@ class Transformer(SequenceToSequence):
 
     def auto_config(self, num_replicas=1):
         config = super().auto_config(num_replicas=num_replicas)
-        return merge_dict(
+        return config_util.merge_config(
             config,
             {
                 "params": {
diff --git a/opennmt/runner.py b/opennmt/runner.py
index c21c79dc..97a0f890 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -114,8 +114,8 @@ class Runner(object):
                 raise NotImplementedError(
                     "This model does not define any automatic configuration values"
                 )
-            misc.merge_dict(config, model_config)
-        misc.merge_dict(config, self._config)
+            config_util.merge_config(config, model_config)
+        config_util.merge_config(config, self._config)
 
         config["params"].setdefault("num_hypotheses", config["infer"].get("n_best", 1))
         config["params"].setdefault(
diff --git a/opennmt/utils/misc.py b/opennmt/utils/misc.py
index c98923e9..4fa96389 100644
--- a/opennmt/utils/misc.py
+++ b/opennmt/utils/misc.py
@@ -384,19 +384,25 @@ def extract_suffixed_keys(dictionary, suffix):
     return sub_dict
 
 
-def merge_dict(dict1, dict2):
+def merge_dict(dict1, dict2, override_keys=None):
     """Merges :obj:`dict2` into :obj:`dict1`.
 
     Args:
       dict1: The base dictionary.
       dict2: The dictionary to merge.
+      override_keys: The values associated with these keys are overridden instead
+        of merged.
 
     Returns:
       The merged dictionary :obj:`dict1`.
     """
     for key, value in dict2.items():
-        if isinstance(value, dict):
-            dict1[key] = merge_dict(dict1.get(key, {}), value)
+        if isinstance(value, dict) and (
+            override_keys is None or key not in override_keys
+        ):
+            dict1[key] = merge_dict(
+                dict1.get(key, {}), value, override_keys=override_keys
+            )
         else:
             dict1[key] = value
     return dict1


Test patch:
diff --git a/opennmt/tests/config_test.py b/opennmt/tests/config_test.py
index ef0a2801..ce88aa73 100644
--- a/opennmt/tests/config_test.py
+++ b/opennmt/tests/config_test.py
@@ -13,8 +13,19 @@ from opennmt.models.model import Model
 
 class ConfigTest(tf.test.TestCase):
     def testConfigOverride(self):
-        config1 = {"model_dir": "foo", "train": {"batch_size": 32, "steps": 42}}
-        config2 = {"model_dir": "bar", "train": {"batch_size": 64}}
+        config1 = {
+            "model_dir": "foo",
+            "train": {"batch_size": 32, "steps": 42},
+            "params": {
+                "optimizer": "Adam",
+                "optimizer_params": {"beta_1": 0.9, "beta_2": 0.998},
+            },
+        }
+        config2 = {
+            "model_dir": "bar",
+            "train": {"batch_size": 64},
+            "params": {"optimizer": "SGD", "optimizer_params": {}},
+        }
         config_file_1 = os.path.join(self.get_temp_dir(), "config1.yml")
         config_file_2 = os.path.join(self.get_temp_dir(), "config2.yml")
 
@@ -26,7 +37,11 @@ class ConfigTest(tf.test.TestCase):
         loaded_config = config.load_config([config_file_1, config_file_2])
 
         self.assertDictEqual(
-            {"model_dir": "bar", "train": {"batch_size": 64, "steps": 42}},
+            {
+                "model_dir": "bar",
+                "train": {"batch_size": 64, "steps": 42},
+                "params": {"optimizer": "SGD", "optimizer_params": {}},
+            },
             loaded_config,
         )
 
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
index 62c7de29..a0e027ce 100644
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -10,7 +10,7 @@ import tensorflow as tf
 from parameterized import parameterized, parameterized_class
 
 from opennmt import Runner, decoders, models
-from opennmt.config import MODEL_DESCRIPTION_FILENAME, load_model
+from opennmt.config import MODEL_DESCRIPTION_FILENAME, load_model, merge_config
 from opennmt.tests import test_util
 from opennmt.utils import exporters, misc
 
@@ -50,7 +50,7 @@ class RunnerTest(tf.test.TestCase):
             "target_vocabulary": "en.vocab",
         }
         if base_config is not None:
-            config = misc.merge_dict(config, base_config)
+            config = merge_config(config, base_config)
         model = load_model(model_dir, as_builder=pass_model_builder)
         runner = Runner(model, config)
         return runner