We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Feature request: Log time for events during training
Hi Guillaume and others, thanks a lot for the great work you are doing!

I was thinking if it's possible to add time tracking to the logger during training. I tried to search the forum and google for relevant issues but didn't find anything. I think it would be helpful for everyone. So please tell me what you think about it.

Thank you,
Alex


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `opennmt/bin/main.py`
Content:
```
1 """Main script."""
2 
3 import argparse
4 import logging
5 import os
6 import sys
7 
8 import tensorflow as tf
9 
10 from opennmt import __version__
11 from opennmt.models import catalog
12 from opennmt.runner import Runner
13 from opennmt.config import load_model, load_config
14 from opennmt.utils import exporters
15 
16 
17 _PYTHON_TO_TENSORFLOW_LOGGING_LEVEL = {
18     logging.CRITICAL: 3,
19     logging.ERROR: 2,
20     logging.WARNING: 1,
21     logging.INFO: 0,
22     logging.DEBUG: 0,
23     logging.NOTSET: 0,
24 }
25 
26 
27 def _set_log_level(log_level):
28     tf.get_logger().setLevel(log_level)
29     os.environ["TF_CPP_MIN_LOG_LEVEL"] = str(
30         _PYTHON_TO_TENSORFLOW_LOGGING_LEVEL[log_level]
31     )
32 
33 
34 def _prefix_paths(prefix, paths):
35     """Recursively prefix paths.
36 
37     Args:
38       prefix: The prefix to apply.
39       data: A dict of relative paths.
40 
41     Returns:
42       The updated dict.
43     """
44     if isinstance(paths, dict):
45         for key, path in paths.items():
46             paths[key] = _prefix_paths(prefix, path)
47         return paths
48     elif isinstance(paths, list):
49         for i, path in enumerate(paths):
50             paths[i] = _prefix_paths(prefix, path)
51         return paths
52     elif isinstance(paths, str):
53         path = paths
54         new_path = os.path.join(prefix, path)
55         if tf.io.gfile.exists(new_path):
56             return new_path
57         else:
58             return path
59     else:
60         return paths
61 
62 
63 def main():
64     parser = argparse.ArgumentParser(
65         formatter_class=argparse.ArgumentDefaultsHelpFormatter
66     )
67     parser.add_argument(
68         "-v", "--version", action="version", version="OpenNMT-tf %s" % __version__
69     )
70     parser.add_argument(
71         "--config", required=True, nargs="+", help="List of configuration files."
72     )
73     parser.add_argument(
74         "--auto_config",
75         default=False,
76         action="store_true",
77         help="Enable automatic configuration values.",
78     )
79     parser.add_argument(
80         "--model_type",
81         default="",
82         choices=list(sorted(catalog.list_model_names_from_catalog())),
83         help="Model type from the catalog.",
84     )
85     parser.add_argument("--model", default="", help="Custom model configuration file.")
86     parser.add_argument(
87         "--run_dir",
88         default="",
89         help="If set, model_dir will be created relative to this location.",
90     )
91     parser.add_argument(
92         "--data_dir",
93         default="",
94         help="If set, data files are expected to be relative to this location.",
95     )
96     parser.add_argument(
97         "--checkpoint_path",
98         default=None,
99         help=(
100             "Specific checkpoint or model directory to load "
101             "(when a directory is set, the latest checkpoint is used)."
102         ),
103     )
104     parser.add_argument(
105         "--log_level",
106         default="INFO",
107         choices=["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "NOTSET"],
108         help="Logs verbosity.",
109     )
110     parser.add_argument("--seed", type=int, default=None, help="Random seed.")
111     parser.add_argument(
112         "--gpu_allow_growth",
113         default=False,
114         action="store_true",
115         help="Allocate GPU memory dynamically.",
116     )
117     parser.add_argument(
118         "--intra_op_parallelism_threads",
119         type=int,
120         default=0,
121         help=(
122             "Number of intra op threads (0 means the system picks "
123             "an appropriate number)."
124         ),
125     )
126     parser.add_argument(
127         "--inter_op_parallelism_threads",
128         type=int,
129         default=0,
130         help=(
131             "Number of inter op threads (0 means the system picks "
132             "an appropriate number)."
133         ),
134     )
135     parser.add_argument(
136         "--mixed_precision",
137         default=False,
138         action="store_true",
139         help="Enable mixed precision.",
140     )
141     parser.add_argument(
142         "--eager_execution",
143         default=False,
144         action="store_true",
145         help="Enable TensorFlow eager execution.",
146     )
147 
148     subparsers = parser.add_subparsers(help="Run type.", dest="run_type")
149     subparsers.required = True
150     parser_train = subparsers.add_parser("train", help="Training.")
151     parser_train.add_argument(
152         "--with_eval",
153         default=False,
154         action="store_true",
155         help="Enable automatic evaluation.",
156     )
157     parser_train.add_argument(
158         "--num_gpus",
159         type=int,
160         default=1,
161         help="Number of GPUs to use for in-graph replication.",
162     )
163     parser_train.add_argument(
164         "--horovod",
165         default=False,
166         action="store_true",
167         help="Enable Horovod training mode.",
168     )
169 
170     parser_eval = subparsers.add_parser("eval", help="Evaluation.")
171     parser_eval.add_argument(
172         "--features_file", nargs="+", default=None, help="Input features files."
173     )
174     parser_eval.add_argument("--labels_file", default=None, help="Output labels files.")
175 
176     parser_infer = subparsers.add_parser("infer", help="Inference.")
177     parser_infer.add_argument(
178         "--features_file", nargs="+", required=True, help="Run inference on this file."
179     )
180     parser_infer.add_argument(
181         "--predictions_file",
182         default="",
183         help=(
184             "File used to save predictions. If not set, predictions are printed "
185             "on the standard output."
186         ),
187     )
188     parser_infer.add_argument(
189         "--log_prediction_time",
190         default=False,
191         action="store_true",
192         help="Logs some prediction time metrics.",
193     )
194 
195     parser_export = subparsers.add_parser("export", help="Model export.")
196     parser_export.add_argument(
197         "--output_dir",
198         "--export_dir",
199         required=True,
200         help="The directory of the exported model.",
201     )
202     parser_export.add_argument(
203         "--format",
204         "--export_format",
205         choices=exporters.list_exporters(),
206         default="saved_model",
207         help="Format of the exported model.",
208     )
209 
210     parser_score = subparsers.add_parser("score", help="Scoring.")
211     parser_score.add_argument(
212         "--features_file", nargs="+", required=True, help="Features file."
213     )
214     parser_score.add_argument(
215         "--predictions_file", default=None, help="Predictions to score."
216     )
217 
218     parser_average_checkpoints = subparsers.add_parser(
219         "average_checkpoints", help="Checkpoint averaging."
220     )
221     parser_average_checkpoints.add_argument(
222         "--output_dir",
223         required=True,
224         help="The output directory for the averaged checkpoint.",
225     )
226     parser_average_checkpoints.add_argument(
227         "--max_count",
228         type=int,
229         default=8,
230         help="The maximal number of checkpoints to average.",
231     )
232 
233     parser_update_vocab = subparsers.add_parser(
234         "update_vocab", help="Update model vocabularies in checkpoint."
235     )
236     parser_update_vocab.add_argument(
237         "--output_dir",
238         required=True,
239         help="The output directory for the updated checkpoint.",
240     )
241     parser_update_vocab.add_argument(
242         "--src_vocab", default=None, help="Path to the new source vocabulary."
243     )
244     parser_update_vocab.add_argument(
245         "--tgt_vocab", default=None, help="Path to the new target vocabulary."
246     )
247 
248     # When using an option that takes multiple values just before the run type,
249     # the run type is treated as a value of this option. To fix this issue, we
250     # inject a placeholder option just before the run type to clearly separate it.
251     parser.add_argument("--placeholder", action="store_true", help=argparse.SUPPRESS)
252     run_types = set(subparsers.choices.keys())
253     args = sys.argv[1:]
254     for i, arg in enumerate(args):
255         if arg in run_types:
256             args.insert(i, "--placeholder")
257             break
258 
259     args = parser.parse_args(args)
260     if (
261         hasattr(args, "features_file")
262         and args.features_file
263         and len(args.features_file) == 1
264     ):
265         args.features_file = args.features_file[0]
266 
267     _set_log_level(getattr(logging, args.log_level))
268     tf.config.threading.set_intra_op_parallelism_threads(
269         args.intra_op_parallelism_threads
270     )
271     tf.config.threading.set_inter_op_parallelism_threads(
272         args.inter_op_parallelism_threads
273     )
274 
275     if args.eager_execution:
276         tf.config.run_functions_eagerly(True)
277 
278     gpus = tf.config.list_physical_devices(device_type="GPU")
279     if hasattr(args, "horovod") and args.horovod:
280         import horovod.tensorflow as hvd
281 
282         hvd.init()
283         is_master = hvd.rank() == 0
284         if gpus:
285             local_gpu = gpus[hvd.local_rank()]
286             tf.config.set_visible_devices(local_gpu, device_type="GPU")
287             gpus = [local_gpu]
288     else:
289         hvd = None
290         is_master = True
291 
292     if args.gpu_allow_growth:
293         for device in gpus:
294             tf.config.experimental.set_memory_growth(device, enable=True)
295 
296     # Load and merge run configurations.
297     config = load_config(args.config)
298     if args.run_dir:
299         config["model_dir"] = os.path.join(args.run_dir, config["model_dir"])
300     if args.data_dir:
301         config["data"] = _prefix_paths(args.data_dir, config["data"])
302 
303     if is_master and not tf.io.gfile.exists(config["model_dir"]):
304         tf.get_logger().info("Creating model directory %s", config["model_dir"])
305         tf.io.gfile.makedirs(config["model_dir"])
306 
307     model = load_model(
308         config["model_dir"],
309         model_file=args.model,
310         model_name=args.model_type,
311         serialize_model=is_master,
312         as_builder=True,
313     )
314     runner = Runner(
315         model,
316         config,
317         auto_config=args.auto_config,
318         mixed_precision=args.mixed_precision,
319         seed=args.seed,
320     )
321 
322     if args.run_type == "train":
323         runner.train(
324             num_devices=args.num_gpus,
325             with_eval=args.with_eval,
326             checkpoint_path=args.checkpoint_path,
327             hvd=hvd,
328         )
329     elif args.run_type == "eval":
330         metrics = runner.evaluate(
331             checkpoint_path=args.checkpoint_path,
332             features_file=args.features_file,
333             labels_file=args.labels_file,
334         )
335         print(metrics)
336     elif args.run_type == "infer":
337         runner.infer(
338             args.features_file,
339             predictions_file=args.predictions_file,
340             checkpoint_path=args.checkpoint_path,
341             log_time=args.log_prediction_time,
342         )
343     elif args.run_type == "export":
344         runner.export(
345             args.output_dir,
346             checkpoint_path=args.checkpoint_path,
347             exporter=exporters.make_exporter(args.format),
348         )
349     elif args.run_type == "score":
350         runner.score(
351             args.features_file,
352             args.predictions_file,
353             checkpoint_path=args.checkpoint_path,
354         )
355     elif args.run_type == "average_checkpoints":
356         runner.average_checkpoints(args.output_dir, max_count=args.max_count)
357     elif args.run_type == "update_vocab":
358         runner.update_vocab(
359             args.output_dir, src_vocab=args.src_vocab, tgt_vocab=args.tgt_vocab
360         )
361 
362 
363 if __name__ == "__main__":
364     main()
365
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/opennmt/bin/main.py b/opennmt/bin/main.py
--- a/opennmt/bin/main.py
+++ b/opennmt/bin/main.py
@@ -24,8 +24,22 @@
 }
 
 
-def _set_log_level(log_level):
-    tf.get_logger().setLevel(log_level)
+def _initialize_logging(log_level):
+    logger = tf.get_logger()
+    logger.setLevel(log_level)
+
+    # Configure the TensorFlow logger to use the same log format as the TensorFlow C++ logs.
+    for handler in list(logger.handlers):
+        logger.removeHandler(handler)
+    formatter = logging.Formatter(
+        fmt="%(asctime)s.%(msecs)03d000: %(levelname).1s %(filename)s:%(lineno)d] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+    # Align the TensorFlow C++ log level with the Python level.
     os.environ["TF_CPP_MIN_LOG_LEVEL"] = str(
         _PYTHON_TO_TENSORFLOW_LOGGING_LEVEL[log_level]
     )
@@ -264,7 +278,7 @@
     ):
         args.features_file = args.features_file[0]
 
-    _set_log_level(getattr(logging, args.log_level))
+    _initialize_logging(getattr(logging, args.log_level))
     tf.config.threading.set_intra_op_parallelism_threads(
         args.intra_op_parallelism_threads
     )

{"golden_diff": "diff --git a/opennmt/bin/main.py b/opennmt/bin/main.py\n--- a/opennmt/bin/main.py\n+++ b/opennmt/bin/main.py\n@@ -24,8 +24,22 @@\n }\n \n \n-def _set_log_level(log_level):\n-    tf.get_logger().setLevel(log_level)\n+def _initialize_logging(log_level):\n+    logger = tf.get_logger()\n+    logger.setLevel(log_level)\n+\n+    # Configure the TensorFlow logger to use the same log format as the TensorFlow C++ logs.\n+    for handler in list(logger.handlers):\n+        logger.removeHandler(handler)\n+    formatter = logging.Formatter(\n+        fmt=\"%(asctime)s.%(msecs)03d000: %(levelname).1s %(filename)s:%(lineno)d] %(message)s\",\n+        datefmt=\"%Y-%m-%d %H:%M:%S\",\n+    )\n+    handler = logging.StreamHandler()\n+    handler.setFormatter(formatter)\n+    logger.addHandler(handler)\n+\n+    # Align the TensorFlow C++ log level with the Python level.\n     os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = str(\n         _PYTHON_TO_TENSORFLOW_LOGGING_LEVEL[log_level]\n     )\n@@ -264,7 +278,7 @@\n     ):\n         args.features_file = args.features_file[0]\n \n-    _set_log_level(getattr(logging, args.log_level))\n+    _initialize_logging(getattr(logging, args.log_level))\n     tf.config.threading.set_intra_op_parallelism_threads(\n         args.intra_op_parallelism_threads\n     )\n", "issue": "Feature request: Log time for events during training\nHi Guillaume and others, thanks a lot for the great work you are doing!\r\n\r\nI was thinking if it's possible to add time tracking to the logger during training. I tried to search the forum and google for relevant issues but didn't find anything. I think it would be helpful for everyone. So please tell me what you think about it.\r\n\r\nThank you,\r\nAlex\r\n\n", "before_files": [{"content": "\"\"\"Main script.\"\"\"\n\nimport argparse\nimport logging\nimport os\nimport sys\n\nimport tensorflow as tf\n\nfrom opennmt import __version__\nfrom opennmt.models import catalog\nfrom opennmt.runner import Runner\nfrom opennmt.config import load_model, load_config\nfrom opennmt.utils import exporters\n\n\n_PYTHON_TO_TENSORFLOW_LOGGING_LEVEL = {\n    logging.CRITICAL: 3,\n    logging.ERROR: 2,\n    logging.WARNING: 1,\n    logging.INFO: 0,\n    logging.DEBUG: 0,\n    logging.NOTSET: 0,\n}\n\n\ndef _set_log_level(log_level):\n    tf.get_logger().setLevel(log_level)\n    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = str(\n        _PYTHON_TO_TENSORFLOW_LOGGING_LEVEL[log_level]\n    )\n\n\ndef _prefix_paths(prefix, paths):\n    \"\"\"Recursively prefix paths.\n\n    Args:\n      prefix: The prefix to apply.\n      data: A dict of relative paths.\n\n    Returns:\n      The updated dict.\n    \"\"\"\n    if isinstance(paths, dict):\n        for key, path in paths.items():\n            paths[key] = _prefix_paths(prefix, path)\n        return paths\n    elif isinstance(paths, list):\n        for i, path in enumerate(paths):\n            paths[i] = _prefix_paths(prefix, path)\n        return paths\n    elif isinstance(paths, str):\n        path = paths\n        new_path = os.path.join(prefix, path)\n        if tf.io.gfile.exists(new_path):\n            return new_path\n        else:\n            return path\n    else:\n        return paths\n\n\ndef main():\n    parser = argparse.ArgumentParser(\n        formatter_class=argparse.ArgumentDefaultsHelpFormatter\n    )\n    parser.add_argument(\n        \"-v\", \"--version\", action=\"version\", version=\"OpenNMT-tf %s\" % __version__\n    )\n    parser.add_argument(\n        \"--config\", required=True, nargs=\"+\", help=\"List of configuration files.\"\n    )\n    parser.add_argument(\n        \"--auto_config\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable automatic configuration values.\",\n    )\n    parser.add_argument(\n        \"--model_type\",\n        default=\"\",\n        choices=list(sorted(catalog.list_model_names_from_catalog())),\n        help=\"Model type from the catalog.\",\n    )\n    parser.add_argument(\"--model\", default=\"\", help=\"Custom model configuration file.\")\n    parser.add_argument(\n        \"--run_dir\",\n        default=\"\",\n        help=\"If set, model_dir will be created relative to this location.\",\n    )\n    parser.add_argument(\n        \"--data_dir\",\n        default=\"\",\n        help=\"If set, data files are expected to be relative to this location.\",\n    )\n    parser.add_argument(\n        \"--checkpoint_path\",\n        default=None,\n        help=(\n            \"Specific checkpoint or model directory to load \"\n            \"(when a directory is set, the latest checkpoint is used).\"\n        ),\n    )\n    parser.add_argument(\n        \"--log_level\",\n        default=\"INFO\",\n        choices=[\"CRITICAL\", \"ERROR\", \"WARNING\", \"INFO\", \"DEBUG\", \"NOTSET\"],\n        help=\"Logs verbosity.\",\n    )\n    parser.add_argument(\"--seed\", type=int, default=None, help=\"Random seed.\")\n    parser.add_argument(\n        \"--gpu_allow_growth\",\n        default=False,\n        action=\"store_true\",\n        help=\"Allocate GPU memory dynamically.\",\n    )\n    parser.add_argument(\n        \"--intra_op_parallelism_threads\",\n        type=int,\n        default=0,\n        help=(\n            \"Number of intra op threads (0 means the system picks \"\n            \"an appropriate number).\"\n        ),\n    )\n    parser.add_argument(\n        \"--inter_op_parallelism_threads\",\n        type=int,\n        default=0,\n        help=(\n            \"Number of inter op threads (0 means the system picks \"\n            \"an appropriate number).\"\n        ),\n    )\n    parser.add_argument(\n        \"--mixed_precision\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable mixed precision.\",\n    )\n    parser.add_argument(\n        \"--eager_execution\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable TensorFlow eager execution.\",\n    )\n\n    subparsers = parser.add_subparsers(help=\"Run type.\", dest=\"run_type\")\n    subparsers.required = True\n    parser_train = subparsers.add_parser(\"train\", help=\"Training.\")\n    parser_train.add_argument(\n        \"--with_eval\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable automatic evaluation.\",\n    )\n    parser_train.add_argument(\n        \"--num_gpus\",\n        type=int,\n        default=1,\n        help=\"Number of GPUs to use for in-graph replication.\",\n    )\n    parser_train.add_argument(\n        \"--horovod\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable Horovod training mode.\",\n    )\n\n    parser_eval = subparsers.add_parser(\"eval\", help=\"Evaluation.\")\n    parser_eval.add_argument(\n        \"--features_file\", nargs=\"+\", default=None, help=\"Input features files.\"\n    )\n    parser_eval.add_argument(\"--labels_file\", default=None, help=\"Output labels files.\")\n\n    parser_infer = subparsers.add_parser(\"infer\", help=\"Inference.\")\n    parser_infer.add_argument(\n        \"--features_file\", nargs=\"+\", required=True, help=\"Run inference on this file.\"\n    )\n    parser_infer.add_argument(\n        \"--predictions_file\",\n        default=\"\",\n        help=(\n            \"File used to save predictions. If not set, predictions are printed \"\n            \"on the standard output.\"\n        ),\n    )\n    parser_infer.add_argument(\n        \"--log_prediction_time\",\n        default=False,\n        action=\"store_true\",\n        help=\"Logs some prediction time metrics.\",\n    )\n\n    parser_export = subparsers.add_parser(\"export\", help=\"Model export.\")\n    parser_export.add_argument(\n        \"--output_dir\",\n        \"--export_dir\",\n        required=True,\n        help=\"The directory of the exported model.\",\n    )\n    parser_export.add_argument(\n        \"--format\",\n        \"--export_format\",\n        choices=exporters.list_exporters(),\n        default=\"saved_model\",\n        help=\"Format of the exported model.\",\n    )\n\n    parser_score = subparsers.add_parser(\"score\", help=\"Scoring.\")\n    parser_score.add_argument(\n        \"--features_file\", nargs=\"+\", required=True, help=\"Features file.\"\n    )\n    parser_score.add_argument(\n        \"--predictions_file\", default=None, help=\"Predictions to score.\"\n    )\n\n    parser_average_checkpoints = subparsers.add_parser(\n        \"average_checkpoints\", help=\"Checkpoint averaging.\"\n    )\n    parser_average_checkpoints.add_argument(\n        \"--output_dir\",\n        required=True,\n        help=\"The output directory for the averaged checkpoint.\",\n    )\n    parser_average_checkpoints.add_argument(\n        \"--max_count\",\n        type=int,\n        default=8,\n        help=\"The maximal number of checkpoints to average.\",\n    )\n\n    parser_update_vocab = subparsers.add_parser(\n        \"update_vocab\", help=\"Update model vocabularies in checkpoint.\"\n    )\n    parser_update_vocab.add_argument(\n        \"--output_dir\",\n        required=True,\n        help=\"The output directory for the updated checkpoint.\",\n    )\n    parser_update_vocab.add_argument(\n        \"--src_vocab\", default=None, help=\"Path to the new source vocabulary.\"\n    )\n    parser_update_vocab.add_argument(\n        \"--tgt_vocab\", default=None, help=\"Path to the new target vocabulary.\"\n    )\n\n    # When using an option that takes multiple values just before the run type,\n    # the run type is treated as a value of this option. To fix this issue, we\n    # inject a placeholder option just before the run type to clearly separate it.\n    parser.add_argument(\"--placeholder\", action=\"store_true\", help=argparse.SUPPRESS)\n    run_types = set(subparsers.choices.keys())\n    args = sys.argv[1:]\n    for i, arg in enumerate(args):\n        if arg in run_types:\n            args.insert(i, \"--placeholder\")\n            break\n\n    args = parser.parse_args(args)\n    if (\n        hasattr(args, \"features_file\")\n        and args.features_file\n        and len(args.features_file) == 1\n    ):\n        args.features_file = args.features_file[0]\n\n    _set_log_level(getattr(logging, args.log_level))\n    tf.config.threading.set_intra_op_parallelism_threads(\n        args.intra_op_parallelism_threads\n    )\n    tf.config.threading.set_inter_op_parallelism_threads(\n        args.inter_op_parallelism_threads\n    )\n\n    if args.eager_execution:\n        tf.config.run_functions_eagerly(True)\n\n    gpus = tf.config.list_physical_devices(device_type=\"GPU\")\n    if hasattr(args, \"horovod\") and args.horovod:\n        import horovod.tensorflow as hvd\n\n        hvd.init()\n        is_master = hvd.rank() == 0\n        if gpus:\n            local_gpu = gpus[hvd.local_rank()]\n            tf.config.set_visible_devices(local_gpu, device_type=\"GPU\")\n            gpus = [local_gpu]\n    else:\n        hvd = None\n        is_master = True\n\n    if args.gpu_allow_growth:\n        for device in gpus:\n            tf.config.experimental.set_memory_growth(device, enable=True)\n\n    # Load and merge run configurations.\n    config = load_config(args.config)\n    if args.run_dir:\n        config[\"model_dir\"] = os.path.join(args.run_dir, config[\"model_dir\"])\n    if args.data_dir:\n        config[\"data\"] = _prefix_paths(args.data_dir, config[\"data\"])\n\n    if is_master and not tf.io.gfile.exists(config[\"model_dir\"]):\n        tf.get_logger().info(\"Creating model directory %s\", config[\"model_dir\"])\n        tf.io.gfile.makedirs(config[\"model_dir\"])\n\n    model = load_model(\n        config[\"model_dir\"],\n        model_file=args.model,\n        model_name=args.model_type,\n        serialize_model=is_master,\n        as_builder=True,\n    )\n    runner = Runner(\n        model,\n        config,\n        auto_config=args.auto_config,\n        mixed_precision=args.mixed_precision,\n        seed=args.seed,\n    )\n\n    if args.run_type == \"train\":\n        runner.train(\n            num_devices=args.num_gpus,\n            with_eval=args.with_eval,\n            checkpoint_path=args.checkpoint_path,\n            hvd=hvd,\n        )\n    elif args.run_type == \"eval\":\n        metrics = runner.evaluate(\n            checkpoint_path=args.checkpoint_path,\n            features_file=args.features_file,\n            labels_file=args.labels_file,\n        )\n        print(metrics)\n    elif args.run_type == \"infer\":\n        runner.infer(\n            args.features_file,\n            predictions_file=args.predictions_file,\n            checkpoint_path=args.checkpoint_path,\n            log_time=args.log_prediction_time,\n        )\n    elif args.run_type == \"export\":\n        runner.export(\n            args.output_dir,\n            checkpoint_path=args.checkpoint_path,\n            exporter=exporters.make_exporter(args.format),\n        )\n    elif args.run_type == \"score\":\n        runner.score(\n            args.features_file,\n            args.predictions_file,\n            checkpoint_path=args.checkpoint_path,\n        )\n    elif args.run_type == \"average_checkpoints\":\n        runner.average_checkpoints(args.output_dir, max_count=args.max_count)\n    elif args.run_type == \"update_vocab\":\n        runner.update_vocab(\n            args.output_dir, src_vocab=args.src_vocab, tgt_vocab=args.tgt_vocab\n        )\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "opennmt/bin/main.py"}], "after_files": [{"content": "\"\"\"Main script.\"\"\"\n\nimport argparse\nimport logging\nimport os\nimport sys\n\nimport tensorflow as tf\n\nfrom opennmt import __version__\nfrom opennmt.models import catalog\nfrom opennmt.runner import Runner\nfrom opennmt.config import load_model, load_config\nfrom opennmt.utils import exporters\n\n\n_PYTHON_TO_TENSORFLOW_LOGGING_LEVEL = {\n    logging.CRITICAL: 3,\n    logging.ERROR: 2,\n    logging.WARNING: 1,\n    logging.INFO: 0,\n    logging.DEBUG: 0,\n    logging.NOTSET: 0,\n}\n\n\ndef _initialize_logging(log_level):\n    logger = tf.get_logger()\n    logger.setLevel(log_level)\n\n    # Configure the TensorFlow logger to use the same log format as the TensorFlow C++ logs.\n    for handler in list(logger.handlers):\n        logger.removeHandler(handler)\n    formatter = logging.Formatter(\n        fmt=\"%(asctime)s.%(msecs)03d000: %(levelname).1s %(filename)s:%(lineno)d] %(message)s\",\n        datefmt=\"%Y-%m-%d %H:%M:%S\",\n    )\n    handler = logging.StreamHandler()\n    handler.setFormatter(formatter)\n    logger.addHandler(handler)\n\n    # Align the TensorFlow C++ log level with the Python level.\n    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = str(\n        _PYTHON_TO_TENSORFLOW_LOGGING_LEVEL[log_level]\n    )\n\n\ndef _prefix_paths(prefix, paths):\n    \"\"\"Recursively prefix paths.\n\n    Args:\n      prefix: The prefix to apply.\n      data: A dict of relative paths.\n\n    Returns:\n      The updated dict.\n    \"\"\"\n    if isinstance(paths, dict):\n        for key, path in paths.items():\n            paths[key] = _prefix_paths(prefix, path)\n        return paths\n    elif isinstance(paths, list):\n        for i, path in enumerate(paths):\n            paths[i] = _prefix_paths(prefix, path)\n        return paths\n    elif isinstance(paths, str):\n        path = paths\n        new_path = os.path.join(prefix, path)\n        if tf.io.gfile.exists(new_path):\n            return new_path\n        else:\n            return path\n    else:\n        return paths\n\n\ndef main():\n    parser = argparse.ArgumentParser(\n        formatter_class=argparse.ArgumentDefaultsHelpFormatter\n    )\n    parser.add_argument(\n        \"-v\", \"--version\", action=\"version\", version=\"OpenNMT-tf %s\" % __version__\n    )\n    parser.add_argument(\n        \"--config\", required=True, nargs=\"+\", help=\"List of configuration files.\"\n    )\n    parser.add_argument(\n        \"--auto_config\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable automatic configuration values.\",\n    )\n    parser.add_argument(\n        \"--model_type\",\n        default=\"\",\n        choices=list(sorted(catalog.list_model_names_from_catalog())),\n        help=\"Model type from the catalog.\",\n    )\n    parser.add_argument(\"--model\", default=\"\", help=\"Custom model configuration file.\")\n    parser.add_argument(\n        \"--run_dir\",\n        default=\"\",\n        help=\"If set, model_dir will be created relative to this location.\",\n    )\n    parser.add_argument(\n        \"--data_dir\",\n        default=\"\",\n        help=\"If set, data files are expected to be relative to this location.\",\n    )\n    parser.add_argument(\n        \"--checkpoint_path\",\n        default=None,\n        help=(\n            \"Specific checkpoint or model directory to load \"\n            \"(when a directory is set, the latest checkpoint is used).\"\n        ),\n    )\n    parser.add_argument(\n        \"--log_level\",\n        default=\"INFO\",\n        choices=[\"CRITICAL\", \"ERROR\", \"WARNING\", \"INFO\", \"DEBUG\", \"NOTSET\"],\n        help=\"Logs verbosity.\",\n    )\n    parser.add_argument(\"--seed\", type=int, default=None, help=\"Random seed.\")\n    parser.add_argument(\n        \"--gpu_allow_growth\",\n        default=False,\n        action=\"store_true\",\n        help=\"Allocate GPU memory dynamically.\",\n    )\n    parser.add_argument(\n        \"--intra_op_parallelism_threads\",\n        type=int,\n        default=0,\n        help=(\n            \"Number of intra op threads (0 means the system picks \"\n            \"an appropriate number).\"\n        ),\n    )\n    parser.add_argument(\n        \"--inter_op_parallelism_threads\",\n        type=int,\n        default=0,\n        help=(\n            \"Number of inter op threads (0 means the system picks \"\n            \"an appropriate number).\"\n        ),\n    )\n    parser.add_argument(\n        \"--mixed_precision\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable mixed precision.\",\n    )\n    parser.add_argument(\n        \"--eager_execution\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable TensorFlow eager execution.\",\n    )\n\n    subparsers = parser.add_subparsers(help=\"Run type.\", dest=\"run_type\")\n    subparsers.required = True\n    parser_train = subparsers.add_parser(\"train\", help=\"Training.\")\n    parser_train.add_argument(\n        \"--with_eval\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable automatic evaluation.\",\n    )\n    parser_train.add_argument(\n        \"--num_gpus\",\n        type=int,\n        default=1,\n        help=\"Number of GPUs to use for in-graph replication.\",\n    )\n    parser_train.add_argument(\n        \"--horovod\",\n        default=False,\n        action=\"store_true\",\n        help=\"Enable Horovod training mode.\",\n    )\n\n    parser_eval = subparsers.add_parser(\"eval\", help=\"Evaluation.\")\n    parser_eval.add_argument(\n        \"--features_file\", nargs=\"+\", default=None, help=\"Input features files.\"\n    )\n    parser_eval.add_argument(\"--labels_file\", default=None, help=\"Output labels files.\")\n\n    parser_infer = subparsers.add_parser(\"infer\", help=\"Inference.\")\n    parser_infer.add_argument(\n        \"--features_file\", nargs=\"+\", required=True, help=\"Run inference on this file.\"\n    )\n    parser_infer.add_argument(\n        \"--predictions_file\",\n        default=\"\",\n        help=(\n            \"File used to save predictions. If not set, predictions are printed \"\n            \"on the standard output.\"\n        ),\n    )\n    parser_infer.add_argument(\n        \"--log_prediction_time\",\n        default=False,\n        action=\"store_true\",\n        help=\"Logs some prediction time metrics.\",\n    )\n\n    parser_export = subparsers.add_parser(\"export\", help=\"Model export.\")\n    parser_export.add_argument(\n        \"--output_dir\",\n        \"--export_dir\",\n        required=True,\n        help=\"The directory of the exported model.\",\n    )\n    parser_export.add_argument(\n        \"--format\",\n        \"--export_format\",\n        choices=exporters.list_exporters(),\n        default=\"saved_model\",\n        help=\"Format of the exported model.\",\n    )\n\n    parser_score = subparsers.add_parser(\"score\", help=\"Scoring.\")\n    parser_score.add_argument(\n        \"--features_file\", nargs=\"+\", required=True, help=\"Features file.\"\n    )\n    parser_score.add_argument(\n        \"--predictions_file\", default=None, help=\"Predictions to score.\"\n    )\n\n    parser_average_checkpoints = subparsers.add_parser(\n        \"average_checkpoints\", help=\"Checkpoint averaging.\"\n    )\n    parser_average_checkpoints.add_argument(\n        \"--output_dir\",\n        required=True,\n        help=\"The output directory for the averaged checkpoint.\",\n    )\n    parser_average_checkpoints.add_argument(\n        \"--max_count\",\n        type=int,\n        default=8,\n        help=\"The maximal number of checkpoints to average.\",\n    )\n\n    parser_update_vocab = subparsers.add_parser(\n        \"update_vocab\", help=\"Update model vocabularies in checkpoint.\"\n    )\n    parser_update_vocab.add_argument(\n        \"--output_dir\",\n        required=True,\n        help=\"The output directory for the updated checkpoint.\",\n    )\n    parser_update_vocab.add_argument(\n        \"--src_vocab\", default=None, help=\"Path to the new source vocabulary.\"\n    )\n    parser_update_vocab.add_argument(\n        \"--tgt_vocab\", default=None, help=\"Path to the new target vocabulary.\"\n    )\n\n    # When using an option that takes multiple values just before the run type,\n    # the run type is treated as a value of this option. To fix this issue, we\n    # inject a placeholder option just before the run type to clearly separate it.\n    parser.add_argument(\"--placeholder\", action=\"store_true\", help=argparse.SUPPRESS)\n    run_types = set(subparsers.choices.keys())\n    args = sys.argv[1:]\n    for i, arg in enumerate(args):\n        if arg in run_types:\n            args.insert(i, \"--placeholder\")\n            break\n\n    args = parser.parse_args(args)\n    if (\n        hasattr(args, \"features_file\")\n        and args.features_file\n        and len(args.features_file) == 1\n    ):\n        args.features_file = args.features_file[0]\n\n    _initialize_logging(getattr(logging, args.log_level))\n    tf.config.threading.set_intra_op_parallelism_threads(\n        args.intra_op_parallelism_threads\n    )\n    tf.config.threading.set_inter_op_parallelism_threads(\n        args.inter_op_parallelism_threads\n    )\n\n    if args.eager_execution:\n        tf.config.run_functions_eagerly(True)\n\n    gpus = tf.config.list_physical_devices(device_type=\"GPU\")\n    if hasattr(args, \"horovod\") and args.horovod:\n        import horovod.tensorflow as hvd\n\n        hvd.init()\n        is_master = hvd.rank() == 0\n        if gpus:\n            local_gpu = gpus[hvd.local_rank()]\n            tf.config.set_visible_devices(local_gpu, device_type=\"GPU\")\n            gpus = [local_gpu]\n    else:\n        hvd = None\n        is_master = True\n\n    if args.gpu_allow_growth:\n        for device in gpus:\n            tf.config.experimental.set_memory_growth(device, enable=True)\n\n    # Load and merge run configurations.\n    config = load_config(args.config)\n    if args.run_dir:\n        config[\"model_dir\"] = os.path.join(args.run_dir, config[\"model_dir\"])\n    if args.data_dir:\n        config[\"data\"] = _prefix_paths(args.data_dir, config[\"data\"])\n\n    if is_master and not tf.io.gfile.exists(config[\"model_dir\"]):\n        tf.get_logger().info(\"Creating model directory %s\", config[\"model_dir\"])\n        tf.io.gfile.makedirs(config[\"model_dir\"])\n\n    model = load_model(\n        config[\"model_dir\"],\n        model_file=args.model,\n        model_name=args.model_type,\n        serialize_model=is_master,\n        as_builder=True,\n    )\n    runner = Runner(\n        model,\n        config,\n        auto_config=args.auto_config,\n        mixed_precision=args.mixed_precision,\n        seed=args.seed,\n    )\n\n    if args.run_type == \"train\":\n        runner.train(\n            num_devices=args.num_gpus,\n            with_eval=args.with_eval,\n            checkpoint_path=args.checkpoint_path,\n            hvd=hvd,\n        )\n    elif args.run_type == \"eval\":\n        metrics = runner.evaluate(\n            checkpoint_path=args.checkpoint_path,\n            features_file=args.features_file,\n            labels_file=args.labels_file,\n        )\n        print(metrics)\n    elif args.run_type == \"infer\":\n        runner.infer(\n            args.features_file,\n            predictions_file=args.predictions_file,\n            checkpoint_path=args.checkpoint_path,\n            log_time=args.log_prediction_time,\n        )\n    elif args.run_type == \"export\":\n        runner.export(\n            args.output_dir,\n            checkpoint_path=args.checkpoint_path,\n            exporter=exporters.make_exporter(args.format),\n        )\n    elif args.run_type == \"score\":\n        runner.score(\n            args.features_file,\n            args.predictions_file,\n            checkpoint_path=args.checkpoint_path,\n        )\n    elif args.run_type == \"average_checkpoints\":\n        runner.average_checkpoints(args.output_dir, max_count=args.max_count)\n    elif args.run_type == \"update_vocab\":\n        runner.update_vocab(\n            args.output_dir, src_vocab=args.src_vocab, tgt_vocab=args.tgt_vocab\n        )\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "opennmt/bin/main.py"}]}