minimal_transformer_training.py

"""This example demonstrates how to train a standard Transformer model in a few
lines of code using OpenNMT-tf high-level APIs.
"""

import argparse
import logging

import tensorflow as tf

import opennmt as onmt

tf.get_logger().setLevel(logging.INFO)


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("run", choices=["train", "translate"], help="Run type.")
    parser.add_argument("--src", required=True, help="Path to the source file.")
    parser.add_argument("--tgt", help="Path to the target file.")
    parser.add_argument(
        "--src_vocab", required=True, help="Path to the source vocabulary."
    )
    parser.add_argument(
        "--tgt_vocab", required=True, help="Path to the target vocabulary."
    )
    parser.add_argument(
        "--model_dir",
        default="checkpoint",
        help="Directory where checkpoint are written.",
    )
    args = parser.parse_args()

    # See http://opennmt.net/OpenNMT-tf/configuration.html for a complete specification
    # of the configuration.
    config = {
        "model_dir": args.model_dir,
        "data": {
            "source_vocabulary": args.src_vocab,
            "target_vocabulary": args.tgt_vocab,
            "train_features_file": args.src,
            "train_labels_file": args.tgt,
        },
    }

    model = onmt.models.TransformerBase()
    runner = onmt.Runner(model, config, auto_config=True)

    if args.run == "train":
        runner.train()
    elif args.run == "translate":
        runner.infer(args.src)


if __name__ == "__main__":
    main()