Repository: OpenNMT/OpenNMT-tf
Pull request #878 (2021-09-24 13:21:15)

Accept data paths relative to the model directory


Patch:
diff --git a/opennmt/bin/main.py b/opennmt/bin/main.py
index 040325ea..f0de2024 100644
--- a/opennmt/bin/main.py
+++ b/opennmt/bin/main.py
@@ -8,7 +8,7 @@ import sys
 import tensorflow as tf
 
 from opennmt import __version__
-from opennmt.config import load_config, load_model
+from opennmt import config as config_util
 from opennmt.models import catalog
 from opennmt.runner import Runner
 from opennmt.utils import exporters
@@ -30,35 +30,6 @@ def _initialize_logging(log_level):
     logger.addHandler(handler)
 
 
-def _prefix_paths(prefix, paths):
-    """Recursively prefix paths.
-
-    Args:
-      prefix: The prefix to apply.
-      data: A dict of relative paths.
-
-    Returns:
-      The updated dict.
-    """
-    if isinstance(paths, dict):
-        for key, path in paths.items():
-            paths[key] = _prefix_paths(prefix, path)
-        return paths
-    elif isinstance(paths, list):
-        for i, path in enumerate(paths):
-            paths[i] = _prefix_paths(prefix, path)
-        return paths
-    elif isinstance(paths, str):
-        path = paths
-        new_path = os.path.join(prefix, path)
-        if tf.io.gfile.exists(new_path):
-            return new_path
-        else:
-            return path
-    else:
-        return paths
-
-
 def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
@@ -294,17 +265,17 @@ def main():
             tf.config.experimental.set_memory_growth(device, enable=True)
 
     # Load and merge run configurations.
-    config = load_config(args.config)
+    config = config_util.load_config(args.config)
     if args.run_dir:
         config["model_dir"] = os.path.join(args.run_dir, config["model_dir"])
     if args.data_dir:
-        config["data"] = _prefix_paths(args.data_dir, config["data"])
+        config["data"] = config_util.try_prefix_paths(args.data_dir, config["data"])
 
     if is_master and not tf.io.gfile.exists(config["model_dir"]):
         tf.get_logger().info("Creating model directory %s", config["model_dir"])
         tf.io.gfile.makedirs(config["model_dir"])
 
-    model = load_model(
+    model = config_util.load_model(
         config["model_dir"],
         model_file=args.model,
         model_name=args.model_type,
diff --git a/opennmt/config.py b/opennmt/config.py
index 2777a296..e8705d49 100644
--- a/opennmt/config.py
+++ b/opennmt/config.py
@@ -157,6 +157,39 @@ def load_config(config_paths, config=None):
     return config
 
 
+def _map_config_values(config, fn):
+    """Applies the function ``fn`` on each value of the configuration."""
+    if isinstance(config, dict):
+        return {key: _map_config_values(value, fn) for key, value in config.items()}
+    elif isinstance(config, list):
+        return [_map_config_values(elem, fn) for elem in config]
+    else:
+        return fn(config)
+
+
+def try_prefix_paths(prefix, config):
+    """Recursively tries to prefix paths in the configuration.
+
+    The path is unchanged if the prefixed path does not exist.
+
+    Args:
+      prefix: The prefix to apply.
+      config: The configuration.
+
+    Returns:
+      The updated configuration.
+    """
+
+    def _fn(path):
+        if isinstance(path, str):
+            new_path = os.path.join(prefix, path)
+            if tf.io.gfile.exists(new_path):
+                return new_path
+        return path
+
+    return _map_config_values(config, _fn)
+
+
 def convert_to_v2_config(v1_config):
     """Converts a V1 configuration to its V2 equivalent.
 
diff --git a/opennmt/runner.py b/opennmt/runner.py
index 41e7dee6..1380c431 100644
--- a/opennmt/runner.py
+++ b/opennmt/runner.py
@@ -13,6 +13,7 @@ import numpy as np
 import tensorflow as tf
 import yaml
 
+from opennmt import config as config_util
 from opennmt import evaluation, inference, models
 from opennmt import training as training_util
 from opennmt.config import MODEL_DESCRIPTION_FILENAME
@@ -82,6 +83,9 @@ class Runner(object):
         tf.get_logger().info("Using model:\n%s", self._model)
         self._optimizer = None
         self._config = copy.deepcopy(config)
+        self._config["data"] = config_util.try_prefix_paths(
+            self._config["model_dir"], config["data"]
+        )
         self._auto_config = auto_config
         self._mixed_precision = mixed_precision
         if seed is not None:


Test patch:
diff --git a/opennmt/tests/config_test.py b/opennmt/tests/config_test.py
index 036c563a..ef0a2801 100644
--- a/opennmt/tests/config_test.py
+++ b/opennmt/tests/config_test.py
@@ -119,6 +119,28 @@ class ConfigTest(tf.test.TestCase):
         with self.assertRaises(RuntimeError):
             config.load_model(self.get_temp_dir())
 
+    def testPathsPrefix(self):
+        a = os.path.join(self.get_temp_dir(), "a.txt")
+        b = os.path.join(self.get_temp_dir(), "b.txt")
+        open(a, "w").close()
+        open(b, "w").close()
+
+        original_config = {
+            "a": "a.txt",
+            "b": ["b.txt"],
+            "c": {"a": "a.txt", "d": "d.txt"},
+        }
+        expected_config = {
+            "a": a,
+            "b": [b],
+            "c": {"a": a, "d": "d.txt"},
+        }
+
+        self.assertDictEqual(
+            config.try_prefix_paths(self.get_temp_dir(), original_config),
+            expected_config,
+        )
+
     def testConvertToV2Config(self):
         v1_config = {
             "data": {
diff --git a/opennmt/tests/runner_test.py b/opennmt/tests/runner_test.py
index f20abdad..62c7de29 100644
--- a/opennmt/tests/runner_test.py
+++ b/opennmt/tests/runner_test.py
@@ -46,8 +46,8 @@ class RunnerTest(tf.test.TestCase):
         config = {}
         config["model_dir"] = model_dir
         config["data"] = {
-            "source_vocabulary": os.path.join(model_dir, "ar.vocab"),
-            "target_vocabulary": os.path.join(model_dir, "en.vocab"),
+            "source_vocabulary": "ar.vocab",
+            "target_vocabulary": "en.vocab",
         }
         if base_config is not None:
             config = misc.merge_dict(config, base_config)
@@ -197,6 +197,7 @@ class RunnerTest(tf.test.TestCase):
             os.path.join(self.get_temp_dir(), "vocab.txt"), list(map(str, range(10)))
         )
         config = {
+            "model_dir": os.path.join(self.get_temp_dir(), "model"),
             "data": {
                 "train_features_file": src,
                 "vocabulary": vocab,