Upload 313 files

infer.sh

@@ -0,0 +1,10 @@
+python inference.py \
+--model trained_models/retrosnyhesis_ReactSeq_prompt_model_on_50k_aug100.pt \
+--src ./tmp_data/src.txt \
+--output ./tmp_data/tgt.txt \
+--beam_size 10 \
+--n_best 10 \
+--batch_size 16384 \
+--batch_type tokens \
+--max_length 500 \
+--seed 0

inference.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python
+from onmt.bin.translate import main
+
+if __name__ == "__main__":
+    main()

onmt/__init__.py

@@ -0,0 +1,24 @@
+""" Main entry point of the ONMT library """
+import onmt.inputters
+import onmt.encoders
+import onmt.decoders
+import onmt.models
+import onmt.utils
+import onmt.modules
+import sys
+import onmt.utils.optimizers
+
+onmt.utils.optimizers.Optim = onmt.utils.optimizers.Optimizer
+sys.modules["onmt.Optim"] = onmt.utils.optimizers
+
+# For Flake
+__all__ = [
+    onmt.inputters,
+    onmt.encoders,
+    onmt.decoders,
+    onmt.models,
+    onmt.utils,
+    onmt.modules,
+]
+
+__version__ = "3.4.1"

onmt/bin/average_models.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+import argparse
+import torch
+
+
+def average_models(model_files, fp32=False):
+    vocab = None
+    opt = None
+    avg_model = None
+    avg_generator = None
+
+    for i, model_file in enumerate(model_files):
+        m = torch.load(model_file, map_location="cpu")
+        model_weights = m["model"]
+        generator_weights = m["generator"]
+
+        if fp32:
+            for k, v in model_weights.items():
+                model_weights[k] = v.float()
+            for k, v in generator_weights.items():
+                generator_weights[k] = v.float()
+
+        if i == 0:
+            vocab, opt = m["vocab"], m["opt"]
+            avg_model = model_weights
+            avg_generator = generator_weights
+        else:
+            for k, v in avg_model.items():
+                avg_model[k].mul_(i).add_(model_weights[k]).div_(i + 1)
+
+            for k, v in avg_generator.items():
+                avg_generator[k].mul_(i).add_(generator_weights[k]).div_(i + 1)
+
+    final = {
+        "vocab": vocab,
+        "opt": opt,
+        "optim": None,
+        "generator": avg_generator,
+        "model": avg_model,
+    }
+    return final
+
+
+def main():
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument(
+        "-models", "-m", nargs="+", required=True, help="List of models"
+    )
+    parser.add_argument("-output", "-o", required=True, help="Output file")
+    parser.add_argument(
+        "-fp32", "-f", action="store_true", help="Cast params to float32"
+    )
+    opt = parser.parse_args()
+
+    final = average_models(opt.models, opt.fp32)
+    torch.save(final, opt.output)
+
+
+if __name__ == "__main__":
+    main()

onmt/bin/build_vocab.py

@@ -0,0 +1,287 @@
+#!/usr/bin/env python
+"""Get vocabulary coutings from transformed corpora samples."""
+import os
+import copy
+import multiprocessing as mp
+import pyonmttok
+from functools import partial
+from onmt.utils.logging import init_logger, logger
+from onmt.utils.misc import set_random_seed, check_path
+from onmt.utils.parse import ArgumentParser
+from onmt.opts import dynamic_prepare_opts
+from onmt.inputters.text_corpus import build_corpora_iters, get_corpora
+from onmt.inputters.text_utils import process, append_features_to_text
+from onmt.transforms import make_transforms, get_transforms_cls
+from onmt.constants import CorpusName, CorpusTask
+from collections import Counter
+
+
+MAXBUCKETSIZE = 256000
+
+
+def write_files_from_queues(sample_path, queues):
+    """
+    Standalone process that reads data from
+    queues in order and write to sample files.
+    """
+    os.makedirs(sample_path, exist_ok=True)
+    for c_name in queues.keys():
+        dest_base = os.path.join(sample_path, "{}.{}".format(c_name, CorpusName.SAMPLE))
+        with open(dest_base + ".src", "w", encoding="utf-8") as f_src, open(
+            dest_base + ".tgt", "w", encoding="utf-8"
+        ) as f_tgt:
+            while True:
+                _next = False
+                for q in queues[c_name]:
+                    item = q.get()
+                    if item == "blank":
+                        continue
+                    if item == "break":
+                        _next = True
+                        break
+                    _, src_line, tgt_line = item
+                    f_src.write(src_line + "\n")
+                    f_tgt.write(tgt_line + "\n")
+                if _next:
+                    break
+
+
+def build_sub_vocab(corpora, transforms, opts, n_sample, stride, offset):
+    """Build vocab on (strided) subpart of the data."""
+    sub_counter_src = Counter()
+    sub_counter_tgt = Counter()
+    sub_counter_src_feats = [Counter() for _ in range(opts.n_src_feats)]
+    datasets_iterables = build_corpora_iters(
+        corpora,
+        transforms,
+        opts.data,
+        skip_empty_level=opts.skip_empty_level,
+        stride=stride,
+        offset=offset,
+    )
+    for c_name, c_iter in datasets_iterables.items():
+        for i, item in enumerate(c_iter):
+            maybe_example = process(CorpusTask.TRAIN, [item])
+            if maybe_example is not None:
+                maybe_example = maybe_example[0]
+            else:
+                if opts.dump_samples:
+                    build_sub_vocab.queues[c_name][offset].put("blank")
+                continue
+            src_line, tgt_line = (
+                maybe_example["src"]["src"],
+                maybe_example["tgt"]["tgt"],
+            )
+            sub_counter_src.update(src_line.split(" "))
+            sub_counter_tgt.update(tgt_line.split(" "))
+
+            if "feats" in maybe_example["src"]:
+                src_feats_lines = maybe_example["src"]["feats"]
+                for k in range(opts.n_src_feats):
+                    sub_counter_src_feats[k].update(src_feats_lines[k].split(" "))
+            else:
+                src_feats_lines = []
+
+            if opts.dump_samples:
+                src_pretty_line = append_features_to_text(src_line, src_feats_lines)
+                build_sub_vocab.queues[c_name][offset].put(
+                    (i, src_pretty_line, tgt_line)
+                )
+            if n_sample > 0 and ((i + 1) * stride + offset) >= n_sample:
+                if opts.dump_samples:
+                    build_sub_vocab.queues[c_name][offset].put("break")
+                break
+        if opts.dump_samples:
+            build_sub_vocab.queues[c_name][offset].put("break")
+    return sub_counter_src, sub_counter_tgt, sub_counter_src_feats
+
+
+def init_pool(queues):
+    """Add the queues as attribute of the pooled function."""
+    build_sub_vocab.queues = queues
+
+
+def build_vocab(opts, transforms, n_sample=3):
+    """Build vocabulary from data."""
+
+    if n_sample == -1:
+        logger.info(f"n_sample={n_sample}: Build vocab on full datasets.")
+    elif n_sample > 0:
+        logger.info(f"Build vocab on {n_sample} transformed examples/corpus.")
+    else:
+        raise ValueError(f"n_sample should > 0 or == -1, get {n_sample}.")
+
+    if opts.dump_samples:
+        logger.info(
+            "The samples on which the vocab is built will be "
+            "dumped to disk. It may slow down the process."
+        )
+    corpora = get_corpora(opts, task=CorpusTask.TRAIN)
+    counter_src = Counter()
+    counter_tgt = Counter()
+    counter_src_feats = [Counter() for _ in range(opts.n_src_feats)]
+
+    queues = {
+        c_name: [
+            mp.Queue(opts.vocab_sample_queue_size) for i in range(opts.num_threads)
+        ]
+        for c_name in corpora.keys()
+    }
+    sample_path = os.path.join(os.path.dirname(opts.save_data), CorpusName.SAMPLE)
+    if opts.dump_samples:
+        write_process = mp.Process(
+            target=write_files_from_queues, args=(sample_path, queues), daemon=True
+        )
+        write_process.start()
+    with mp.Pool(opts.num_threads, init_pool, [queues]) as p:
+        func = partial(
+            build_sub_vocab, corpora, transforms, opts, n_sample, opts.num_threads
+        )
+        for sub_counter_src, sub_counter_tgt, sub_counter_src_feats in p.imap(
+            func, range(0, opts.num_threads)
+        ):
+            counter_src.update(sub_counter_src)
+            counter_tgt.update(sub_counter_tgt)
+            for i in range(opts.n_src_feats):
+                counter_src_feats[i].update(sub_counter_src_feats[i])
+    if opts.dump_samples:
+        write_process.join()
+    return counter_src, counter_tgt, counter_src_feats
+
+
+def ingest_tokens(opts, transforms, n_sample, learner, stride, offset):
+    def _mp_ingest(data):
+        func = partial(process, CorpusName.TRAIN)
+        chunk = len(data) // opts.num_threads
+        with mp.Pool(opts.num_threads) as pool:
+            buckets = pool.map(
+                func,
+                [data[i * chunk : (i + 1) * chunk] for i in range(0, opts.num_threads)],
+            )
+        for bucket in buckets:
+            for ex in bucket:
+                if ex is not None:
+                    src_line, tgt_line = (ex["src"]["src"], ex["tgt"]["tgt"])
+                    learner.ingest(src_line)
+                    learner.ingest(tgt_line)
+
+    corpora = get_corpora(opts, task=CorpusTask.TRAIN)
+    datasets_iterables = build_corpora_iters(
+        corpora,
+        transforms,
+        opts.data,
+        skip_empty_level=opts.skip_empty_level,
+        stride=stride,
+        offset=offset,
+    )
+    to_ingest = []
+    for c_name, c_iter in datasets_iterables.items():
+        for i, item in enumerate(c_iter):
+            if n_sample >= 0 and i >= n_sample:
+                break
+            if len(to_ingest) >= MAXBUCKETSIZE:
+                _mp_ingest(to_ingest)
+                to_ingest = []
+            to_ingest.append(item)
+        _mp_ingest(to_ingest)
+
+
+def make_learner(tokenization_type, symbols):
+    if tokenization_type == "bpe":
+        # BPE training
+        learner = pyonmttok.BPELearner(tokenizer=None, symbols=symbols)
+    elif tokenization_type == "sentencepiece":
+        # SentencePiece training
+        learner = pyonmttok.SentencePieceLearner(
+            vocab_size=symbols, character_coverage=0.98
+        )
+    return learner
+
+
+def build_vocab_main(opts):
+    """Apply transforms to samples of specified data and build vocab from it.
+
+    Transforms that need vocab will be disabled in this.
+    Built vocab is saved in plain text format as following and can be pass as
+    `-src_vocab` (and `-tgt_vocab`) when training:
+    ```
+    <tok_0>\t<count_0>
+    <tok_1>\t<count_1>
+    ```
+    """
+
+    ArgumentParser.validate_prepare_opts(opts, build_vocab_only=True)
+    assert (
+        opts.n_sample == -1 or opts.n_sample > 1
+    ), f"Illegal argument n_sample={opts.n_sample}."
+
+    logger = init_logger()
+    set_random_seed(opts.seed, False)
+    transforms_cls = get_transforms_cls(opts._all_transform)
+
+    if opts.learn_subwords:
+        logger.info(f"Ingesting {opts.src_subword_type} model from corpus")
+        learner = make_learner(opts.src_subword_type, opts.learn_subwords_size)
+        if opts.src_subword_model is not None:
+            tok_path = opts.src_subword_model
+        else:
+            data_dir = os.path.split(opts.save_data)[0]
+            if not os.path.exists(data_dir):
+                os.makedirs(data_dir)
+            tok_path = os.path.join(data_dir, f"{opts.src_subword_type}.model")
+        save_opts = copy.deepcopy(opts)
+        opts.src_subword_type = "none"
+        opts.tgt_subword_type = "none"
+        opts.src_onmttok_kwargs["joiner_annotate"] = False
+        opts.tgt_onmttok_kwargs["joiner_annotate"] = False
+        transforms = make_transforms(opts, transforms_cls, None)
+        ingest_tokens(opts, transforms, opts.n_sample, learner, 1, 0)
+        logger.info(f"Learning {tok_path} model, patience")
+        learner.learn(tok_path)
+        opts = save_opts
+
+    transforms = make_transforms(opts, transforms_cls, None)
+
+    logger.info(f"Counter vocab from {opts.n_sample} samples.")
+    src_counter, tgt_counter, src_feats_counter = build_vocab(
+        opts, transforms, n_sample=opts.n_sample
+    )
+
+    logger.info(f"Counters src: {len(src_counter)}")
+    logger.info(f"Counters tgt: {len(tgt_counter)}")
+    for i, feat_counter in enumerate(src_feats_counter):
+        logger.info(f"Counters src feat_{i}: {len(feat_counter)}")
+
+    def save_counter(counter, save_path):
+        check_path(save_path, exist_ok=opts.overwrite, log=logger.warning)
+        with open(save_path, "w", encoding="utf8") as fo:
+            for tok, count in counter.most_common():
+                fo.write(tok + "\t" + str(count) + "\n")
+
+    if opts.share_vocab:
+        src_counter += tgt_counter
+        tgt_counter = src_counter
+        logger.info(f"Counters after share:{len(src_counter)}")
+        save_counter(src_counter, opts.src_vocab)
+    else:
+        save_counter(src_counter, opts.src_vocab)
+        save_counter(tgt_counter, opts.tgt_vocab)
+
+    for i, c in enumerate(src_feats_counter):
+        save_counter(c, f"{opts.src_vocab}_feat{i}")
+
+
+def _get_parser():
+    parser = ArgumentParser(description="build_vocab.py")
+    dynamic_prepare_opts(parser, build_vocab_only=True)
+    return parser
+
+
+def main():
+    parser = _get_parser()
+    opts, unknown = parser.parse_known_args()
+    build_vocab_main(opts)
+
+
+if __name__ == "__main__":
+    main()

onmt/bin/release_model.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+import argparse
+import torch
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Release an OpenNMT-py model for inference"
+    )
+    parser.add_argument("--model", "-m", help="The model path", required=True)
+    parser.add_argument("--output", "-o", help="The output path", required=True)
+    parser.add_argument(
+        "--format",
+        choices=["pytorch", "ctranslate2"],
+        default="pytorch",
+        help="The format of the released model",
+    )
+    parser.add_argument(
+        "--quantization",
+        "-q",
+        choices=["int8", "int16", "float16", "int8_float16"],
+        default=None,
+        help="Quantization type for CT2 model.",
+    )
+    opt = parser.parse_args()
+
+    model = torch.load(opt.model, map_location=torch.device("cpu"))
+    if opt.format == "pytorch":
+        model["optim"] = None
+        torch.save(model, opt.output)
+    elif opt.format == "ctranslate2":
+        import ctranslate2
+
+        converter = ctranslate2.converters.OpenNMTPyConverter(opt.model)
+        converter.convert(opt.output, force=True, quantization=opt.quantization)
+
+
+if __name__ == "__main__":
+    main()

onmt/bin/server.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python
+import configargparse
+
+from flask import Flask, jsonify, request
+from waitress import serve
+from onmt.translate import TranslationServer, ServerModelError
+import logging
+from logging.handlers import RotatingFileHandler
+
+STATUS_OK = "ok"
+STATUS_ERROR = "error"
+
+
+def start(config_file, url_root="./translator", host="0.0.0.0", port=5000, debug=False):
+    def prefix_route(route_function, prefix="", mask="{0}{1}"):
+        def newroute(route, *args, **kwargs):
+            return route_function(mask.format(prefix, route), *args, **kwargs)
+
+        return newroute
+
+    if debug:
+        logger = logging.getLogger("main")
+        log_format = logging.Formatter("[%(asctime)s %(levelname)s] %(message)s")
+        file_handler = RotatingFileHandler(
+            "debug_requests.log", maxBytes=1000000, backupCount=10
+        )
+        file_handler.setFormatter(log_format)
+        logger.addHandler(file_handler)
+
+    app = Flask(__name__)
+    app.route = prefix_route(app.route, url_root)
+    translation_server = TranslationServer()
+    translation_server.start(config_file)
+
+    @app.route("/models", methods=["GET"])
+    def get_models():
+        out = translation_server.list_models()
+        return jsonify(out)
+
+    @app.route("/health", methods=["GET"])
+    def health():
+        out = {}
+        out["status"] = STATUS_OK
+        return jsonify(out)
+
+    @app.route("/clone_model/<int:model_id>", methods=["POST"])
+    def clone_model(model_id):
+        out = {}
+        data = request.get_json(force=True)
+        timeout = -1
+        if "timeout" in data:
+            timeout = data["timeout"]
+            del data["timeout"]
+
+        opt = data.get("opt", None)
+        try:
+            model_id, load_time = translation_server.clone_model(model_id, opt, timeout)
+        except ServerModelError as e:
+            out["status"] = STATUS_ERROR
+            out["error"] = str(e)
+        else:
+            out["status"] = STATUS_OK
+            out["model_id"] = model_id
+            out["load_time"] = load_time
+
+        return jsonify(out)
+
+    @app.route("/unload_model/<int:model_id>", methods=["GET"])
+    def unload_model(model_id):
+        out = {"model_id": model_id}
+
+        try:
+            translation_server.unload_model(model_id)
+            out["status"] = STATUS_OK
+        except Exception as e:
+            out["status"] = STATUS_ERROR
+            out["error"] = str(e)
+
+        return jsonify(out)
+
+    @app.route("/translate", methods=["POST"])
+    def translate():
+        inputs = request.get_json(force=True)
+        if debug:
+            logger.info(inputs)
+        out = {}
+        try:
+            trans, scores, n_best, _, aligns, align_scores = translation_server.run(
+                inputs
+            )
+            assert len(trans) == len(inputs) * n_best
+            assert len(scores) == len(inputs) * n_best
+            assert len(aligns) == len(inputs) * n_best
+
+            out = [[] for _ in range(n_best)]
+            for i in range(len(trans)):
+                response = {
+                    "src": inputs[i // n_best]["src"],
+                    "tgt": trans[i],
+                    "n_best": n_best,
+                    "pred_score": scores[i],
+                }
+                if len(aligns[i]) > 0 and aligns[i][0] is not None:
+                    response["align"] = aligns[i]
+                    response["align_score"] = align_scores[i]
+                out[i % n_best].append(response)
+        except ServerModelError as e:
+            model_id = inputs[0].get("id")
+            if debug:
+                logger.warning(
+                    "Unload model #{} " "because of an error".format(model_id)
+                )
+            translation_server.models[model_id].unload()
+            out["error"] = str(e)
+            out["status"] = STATUS_ERROR
+        if debug:
+            logger.info(out)
+        return jsonify(out)
+
+    @app.route("/to_cpu/<int:model_id>", methods=["GET"])
+    def to_cpu(model_id):
+        out = {"model_id": model_id}
+        translation_server.models[model_id].to_cpu()
+
+        out["status"] = STATUS_OK
+        return jsonify(out)
+
+    @app.route("/to_gpu/<int:model_id>", methods=["GET"])
+    def to_gpu(model_id):
+        out = {"model_id": model_id}
+        translation_server.models[model_id].to_gpu()
+
+        out["status"] = STATUS_OK
+        return jsonify(out)
+
+    serve(app, host=host, port=port)
+
+
+def _get_parser():
+    parser = configargparse.ArgumentParser(
+        config_file_parser_class=configargparse.YAMLConfigFileParser,
+        description="OpenNMT-py REST Server",
+    )
+    parser.add_argument("--ip", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default="5000")
+    parser.add_argument("--url_root", type=str, default="/translator")
+    parser.add_argument("--debug", "-d", action="store_true")
+    parser.add_argument(
+        "--config", "-c", type=str, default="./available_models/conf.json"
+    )
+    return parser
+
+
+def main():
+    parser = _get_parser()
+    args = parser.parse_args()
+    start(
+        args.config,
+        url_root=args.url_root,
+        host=args.ip,
+        port=args.port,
+        debug=args.debug,
+    )
+
+
+if __name__ == "__main__":
+    main()

onmt/bin/train.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+"""Train models with dynamic data."""
+import torch
+from functools import partial
+from onmt.utils.distributed import ErrorHandler, spawned_train
+from onmt.utils.misc import set_random_seed
+from onmt.utils.logging import init_logger, logger
+from onmt.utils.parse import ArgumentParser
+from onmt.opts import train_opts
+from onmt.train_single import main as single_main
+
+
+# Set sharing strategy manually instead of default based on the OS.
+# torch.multiprocessing.set_sharing_strategy('file_system')
+
+
+def train(opt):
+    init_logger(opt.log_file)
+
+    ArgumentParser.validate_train_opts(opt)
+    ArgumentParser.update_model_opts(opt)
+    ArgumentParser.validate_model_opts(opt)
+
+    set_random_seed(opt.seed, False)
+
+    train_process = partial(single_main)
+
+    nb_gpu = len(opt.gpu_ranks)
+
+    if opt.world_size > 1:
+        mp = torch.multiprocessing.get_context("spawn")
+        # Create a thread to listen for errors in the child processes.
+        error_queue = mp.SimpleQueue()
+        error_handler = ErrorHandler(error_queue)
+        # Train with multiprocessing.
+        procs = []
+        for device_id in range(nb_gpu):
+            procs.append(
+                mp.Process(
+                    target=spawned_train,
+                    args=(train_process, opt, device_id, error_queue),
+                    daemon=False,
+                )
+            )
+            procs[device_id].start()
+            logger.info(" Starting process pid: %d  " % procs[device_id].pid)
+            error_handler.add_child(procs[device_id].pid)
+        for p in procs:
+            p.join()
+
+    elif nb_gpu == 1:  # case 1 GPU only
+        train_process(opt, device_id=0)
+    else:  # case only CPU
+        train_process(opt, device_id=-1)
+
+
+def _get_parser():
+    parser = ArgumentParser(description="train.py")
+    train_opts(parser)
+    return parser
+
+
+def main():
+    parser = _get_parser()
+
+    opt, unknown = parser.parse_known_args()
+    train(opt)
+
+
+if __name__ == "__main__":
+    main()

onmt/bin/translate.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+from onmt.utils.logging import init_logger
+from onmt.translate.translator import build_translator
+from onmt.inputters.dynamic_iterator import build_dynamic_dataset_iter
+from onmt.inputters.inputter import IterOnDevice
+from onmt.transforms import get_transforms_cls
+from onmt.constants import CorpusTask
+import onmt.opts as opts
+from onmt.utils.parse import ArgumentParser
+from onmt.utils.misc import use_gpu, set_random_seed
+
+
+def translate(opt):
+    ArgumentParser.validate_translate_opts(opt)
+    ArgumentParser._get_all_transform_translate(opt)
+    ArgumentParser._validate_transforms_opts(opt)
+    ArgumentParser.validate_translate_opts_dynamic(opt)
+    logger = init_logger(opt.log_file)
+
+    set_random_seed(opt.seed, use_gpu(opt))
+
+    translator = build_translator(opt, logger=logger, report_score=True)
+
+    transforms_cls = get_transforms_cls(opt._all_transform)
+
+    infer_iter = build_dynamic_dataset_iter(
+        opt,
+        transforms_cls,
+        translator.vocabs,
+        task=CorpusTask.INFER,
+        copy=translator.copy_attn,
+    )
+
+    infer_iter = IterOnDevice(infer_iter, opt.gpu)
+
+    _, _ = translator._translate(
+        infer_iter,
+        transform=infer_iter.transform,
+        attn_debug=opt.attn_debug,
+        align_debug=opt.align_debug,
+    )
+
+
+def _get_parser():
+    parser = ArgumentParser(description="translate.py")
+
+    opts.config_opts(parser)
+    opts.translate_opts(parser, dynamic=True)
+    return parser
+
+
+def main():
+    parser = _get_parser()
+    opt = parser.parse_args()
+    translate(opt)
+
+
+if __name__ == "__main__":
+    main()

onmt/constants.py

@@ -0,0 +1,41 @@
+"""Define constant values used across the project."""
+
+
+class DefaultTokens(object):
+    PAD = "<blank>"
+    BOS = "<s>"
+    EOS = "</s>"
+    UNK = "<unk>"
+    MASK = "<mask>"
+    VOCAB_PAD = "averyunlikelytoken"
+    SENT_FULL_STOPS = [".", "?", "!"]
+    PHRASE_TABLE_SEPARATOR = "|||"
+    ALIGNMENT_SEPARATOR = " ||| "
+    SEP = "｟newline｠"
+    MASK_BEFORE = "｟_mask_before_｠"
+
+
+class CorpusName(object):
+    VALID = "valid"
+    TRAIN = "train"
+    SAMPLE = "sample"
+    INFER = "infer"
+
+
+class CorpusTask(object):
+    TRAIN = "train"
+    VALID = "valid"
+    INFER = "infer"
+
+
+class SubwordMarker(object):
+    SPACER = "▁"
+    JOINER = "￭"
+    BEGIN_UPPERCASE = "｟mrk_begin_case_region_U｠"
+    END_UPPERCASE = "｟mrk_end_case_region_U｠"
+    BEGIN_CASED = "｟mrk_case_modifier_C｠"
+
+
+class ModelTask(object):
+    LANGUAGE_MODEL = "lm"
+    SEQ2SEQ = "seq2seq"

onmt/decoders/__init__.py

@@ -0,0 +1,63 @@
+"""Module defining decoders."""
+import os
+import importlib
+from onmt.decoders.decoder import DecoderBase, InputFeedRNNDecoder, StdRNNDecoder
+from onmt.decoders.transformer import TransformerDecoder, TransformerLMDecoder
+from onmt.decoders.cnn_decoder import CNNDecoder
+
+
+str2dec = {
+    "rnn": StdRNNDecoder,
+    "ifrnn": InputFeedRNNDecoder,
+    "cnn": CNNDecoder,
+    "transformer": TransformerDecoder,
+    "transformer_lm": TransformerLMDecoder,
+}
+
+__all__ = [
+    "DecoderBase",
+    "TransformerDecoder",
+    "StdRNNDecoder",
+    "CNNDecoder",
+    "InputFeedRNNDecoder",
+    "str2dec",
+    "TransformerLMDecoder",
+]
+
+
+def get_decoders_cls(decoders_names):
+    """Return valid encoder class indicated in `decoders_names`."""
+    decoders_cls = {}
+    for name in decoders_names:
+        if name not in str2dec:
+            raise ValueError("%s decoder not supported!" % name)
+        decoders_cls[name] = str2dec[name]
+    return decoders_cls
+
+
+def register_decoder(name):
+    """Encoder register that can be used to add new encoder class."""
+
+    def register_decoder_cls(cls):
+        if name in str2dec:
+            raise ValueError("Cannot register duplicate decoder ({})".format(name))
+        if not issubclass(cls, DecoderBase):
+            raise ValueError(f"decoder ({name}: {cls.__name_}) must extend DecoderBase")
+        str2dec[name] = cls
+        __all__.append(cls.__name__)  # added to be complete
+        return cls
+
+    return register_decoder_cls
+
+
+# Auto import python files in this directory
+decoder_dir = os.path.dirname(__file__)
+for file in os.listdir(decoder_dir):
+    path = os.path.join(decoder_dir, file)
+    if (
+        not file.startswith("_")
+        and not file.startswith(".")
+        and (file.endswith(".py") or os.path.isdir(path))
+    ):
+        file_name = file[: file.find(".py")] if file.endswith(".py") else file
+        module = importlib.import_module("onmt.decoders." + file_name)

onmt/decoders/cnn_decoder.py

@@ -0,0 +1,141 @@
+"""Implementation of the CNN Decoder part of
+"Convolutional Sequence to Sequence Learning"
+"""
+import torch
+import torch.nn as nn
+
+from onmt.modules import ConvMultiStepAttention, GlobalAttention
+from onmt.utils.cnn_factory import shape_transform, GatedConv
+from onmt.decoders.decoder import DecoderBase
+
+SCALE_WEIGHT = 0.5**0.5
+
+
+class CNNDecoder(DecoderBase):
+    """Decoder based on "Convolutional Sequence to Sequence Learning"
+    :cite:`DBLP:journals/corr/GehringAGYD17`.
+
+    Consists of residual convolutional layers, with ConvMultiStepAttention.
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        hidden_size,
+        attn_type,
+        copy_attn,
+        cnn_kernel_width,
+        dropout,
+        embeddings,
+        copy_attn_type,
+    ):
+        super(CNNDecoder, self).__init__()
+
+        self.cnn_kernel_width = cnn_kernel_width
+        self.embeddings = embeddings
+
+        # Decoder State
+        self.state = {}
+
+        input_size = self.embeddings.embedding_size
+        self.linear = nn.Linear(input_size, hidden_size)
+        self.conv_layers = nn.ModuleList(
+            [
+                GatedConv(hidden_size, cnn_kernel_width, dropout, True)
+                for i in range(num_layers)
+            ]
+        )
+        self.attn_layers = nn.ModuleList(
+            [ConvMultiStepAttention(hidden_size) for i in range(num_layers)]
+        )
+
+        # CNNDecoder has its own attention mechanism.
+        # Set up a separate copy attention layer if needed.
+        assert not copy_attn, "Copy mechanism not yet tested in conv2conv"
+        if copy_attn:
+            self.copy_attn = GlobalAttention(hidden_size, attn_type=copy_attn_type)
+        else:
+            self.copy_attn = None
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor."""
+        return cls(
+            opt.dec_layers,
+            opt.dec_hid_size,
+            opt.global_attention,
+            opt.copy_attn,
+            opt.cnn_kernel_width,
+            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
+            embeddings,
+            opt.copy_attn_type,
+        )
+
+    def init_state(self, _, enc_out, enc_hidden):
+        """Init decoder state."""
+        self.state["src"] = (enc_out + enc_hidden) * SCALE_WEIGHT
+        self.state["previous_input"] = None
+
+    def map_state(self, fn):
+        self.state["src"] = fn(self.state["src"], 0)
+        if self.state["previous_input"] is not None:
+            self.state["previous_input"] = fn(self.state["previous_input"], 0)
+
+    def detach_state(self):
+        self.state["previous_input"] = self.state["previous_input"].detach()
+
+    def forward(self, tgt, enc_out, step=None, **kwargs):
+        """See :obj:`onmt.modules.RNNDecoderBase.forward()`"""
+
+        if self.state["previous_input"] is not None:
+            tgt = torch.cat([self.state["previous_input"], tgt], 1)
+
+        dec_outs = []
+        attns = {"std": []}
+        if self.copy_attn is not None:
+            attns["copy"] = []
+
+        emb = self.embeddings(tgt)
+        assert emb.dim() == 3  # batch x len x embedding_dim
+
+        tgt_emb = emb
+        # The output of CNNEncoder.
+        enc_out_t = enc_out
+        # The combination of output of CNNEncoder and source embeddings.
+        enc_out_c = self.state["src"]
+
+        emb_reshape = tgt_emb.view(tgt_emb.size(0) * tgt_emb.size(1), -1)
+        linear_out = self.linear(emb_reshape)
+        x = linear_out.view(tgt_emb.size(0), tgt_emb.size(1), -1)
+        x = shape_transform(x)
+
+        pad = torch.zeros(x.size(0), x.size(1), self.cnn_kernel_width - 1, 1)
+
+        pad = pad.type_as(x)
+        base_target_emb = x
+
+        for conv, attention in zip(self.conv_layers, self.attn_layers):
+            new_target_input = torch.cat([pad, x], 2)
+            out = conv(new_target_input)
+            c, attn = attention(base_target_emb, out, enc_out_t, enc_out_c)
+            x = (x + (c + out) * SCALE_WEIGHT) * SCALE_WEIGHT
+
+        dec_outs = x.squeeze(3).transpose(1, 2)
+
+        # Process the result and update the attentions.
+        if self.state["previous_input"] is not None:
+            dec_outs = dec_outs[:, self.state["previous_input"].size(1) :, :]
+            attn = attn[:, self.state["previous_input"].size(1) :].squeeze()
+            attn = torch.stack([attn])
+        attns["std"] = attn
+        if self.copy_attn is not None:
+            attns["copy"] = attn
+
+        # Update the state.
+        self.state["previous_input"] = tgt
+        # TODO change the way attns is returned dict => list or tuple (onnx)
+        return dec_outs, attns
+
+    def update_dropout(self, dropout, attention_dropout=None):
+        for layer in self.conv_layers:
+            layer.dropout.p = dropout

onmt/decoders/decoder.py

@@ -0,0 +1,405 @@
+import torch
+import torch.nn as nn
+
+from onmt.modules.stacked_rnn import StackedLSTM, StackedGRU
+from onmt.modules import context_gate_factory, GlobalAttention
+from onmt.utils.rnn_factory import rnn_factory
+
+
+class DecoderBase(nn.Module):
+    """Abstract class for decoders.
+
+    Args:
+        attentional (bool): The decoder returns non-empty attention.
+    """
+
+    def __init__(self, attentional=True):
+        super(DecoderBase, self).__init__()
+        self.attentional = attentional
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor.
+
+        Subclasses should override this method.
+        """
+
+        raise NotImplementedError
+
+
+class RNNDecoderBase(DecoderBase):
+    """Base recurrent attention-based decoder class.
+
+    Specifies the interface used by different decoder types
+    and required by :class:`~onmt.models.NMTModel`.
+
+    Args:
+       rnn_type (str):
+          style of recurrent unit to use, one of [RNN, LSTM, GRU, SRU]
+       bidirectional_encoder (bool) : use with a bidirectional encoder
+       num_layers (int) : number of stacked layers
+       hidden_size (int) : hidden size of each layer
+       attn_type (str) : see :class:`~onmt.modules.GlobalAttention`
+       attn_func (str) : see :class:`~onmt.modules.GlobalAttention`
+       coverage_attn (str): see :class:`~onmt.modules.GlobalAttention`
+       context_gate (str): see :class:`~onmt.modules.ContextGate`
+       copy_attn (bool): setup a separate copy attention mechanism
+       dropout (float) : dropout value for :class:`torch.nn.Dropout`
+       embeddings (onmt.modules.Embeddings): embedding module to use
+       reuse_copy_attn (bool): reuse the attention for copying
+       copy_attn_type (str): The copy attention style. See
+        :class:`~onmt.modules.GlobalAttention`.
+    """
+
+    def __init__(
+        self,
+        rnn_type,
+        bidirectional_encoder,
+        num_layers,
+        hidden_size,
+        attn_type="general",
+        attn_func="softmax",
+        coverage_attn=False,
+        context_gate=None,
+        copy_attn=False,
+        dropout=0.0,
+        embeddings=None,
+        reuse_copy_attn=False,
+        copy_attn_type="general",
+    ):
+        super(RNNDecoderBase, self).__init__(
+            attentional=attn_type != "none" and attn_type is not None
+        )
+
+        self.bidirectional_encoder = bidirectional_encoder
+        self.num_layers = num_layers
+        self.hidden_size = hidden_size
+        self.embeddings = embeddings
+        self.dropout = nn.Dropout(dropout)
+
+        # Decoder state
+        self.state = {}
+
+        # Build the RNN.
+        self.rnn = self._build_rnn(
+            rnn_type,
+            input_size=self._input_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            dropout=dropout,
+        )
+
+        # Set up the context gate.
+        self.context_gate = None
+        if context_gate is not None:
+            self.context_gate = context_gate_factory(
+                context_gate, self._input_size, hidden_size, hidden_size, hidden_size
+            )
+
+        # Set up the standard attention.
+        self._coverage = coverage_attn
+        if not self.attentional:
+            if self._coverage:
+                raise ValueError("Cannot use coverage term with no attention.")
+            self.attn = None
+        else:
+            self.attn = GlobalAttention(
+                hidden_size,
+                coverage=coverage_attn,
+                attn_type=attn_type,
+                attn_func=attn_func,
+            )
+
+        if copy_attn and not reuse_copy_attn:
+            if copy_attn_type == "none" or copy_attn_type is None:
+                raise ValueError("Cannot use copy_attn with copy_attn_type none")
+            self.copy_attn = GlobalAttention(
+                hidden_size, attn_type=copy_attn_type, attn_func=attn_func
+            )
+        else:
+            self.copy_attn = None
+
+        self._reuse_copy_attn = reuse_copy_attn and copy_attn
+        if self._reuse_copy_attn and not self.attentional:
+            raise ValueError("Cannot reuse copy attention with no attention.")
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor."""
+        return cls(
+            opt.rnn_type,
+            opt.brnn,
+            opt.dec_layers,
+            opt.dec_hid_size,
+            opt.global_attention,
+            opt.global_attention_function,
+            opt.coverage_attn,
+            opt.context_gate,
+            opt.copy_attn,
+            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
+            embeddings,
+            opt.reuse_copy_attn,
+            opt.copy_attn_type,
+        )
+
+    def init_state(self, src, _, enc_final_hs):
+        """Initialize decoder state with last state of the encoder."""
+
+        def _fix_enc_hidden(hidden):
+            # The encoder hidden is  (layers*directions) x batch x dim.
+            # We need to convert it to layers x batch x (directions*dim).
+            if self.bidirectional_encoder:
+                hidden = torch.cat(
+                    [hidden[0 : hidden.size(0) : 2], hidden[1 : hidden.size(0) : 2]], 2
+                )
+            return hidden
+
+        if isinstance(enc_final_hs, tuple):  # LSTM
+            self.state["hidden"] = tuple(
+                _fix_enc_hidden(enc_hid) for enc_hid in enc_final_hs
+            )
+        else:  # GRU
+            self.state["hidden"] = (_fix_enc_hidden(enc_final_hs),)
+
+        # Init the input feed.
+        batch_size = self.state["hidden"][0].size(1)
+
+        h_size = (batch_size, self.hidden_size)
+        self.state["input_feed"] = (
+            self.state["hidden"][0].data.new(*h_size).zero_().unsqueeze(0)
+        )
+
+        self.state["coverage"] = None
+
+    def map_state(self, fn):
+        self.state["hidden"] = tuple(fn(h, 1) for h in self.state["hidden"])
+        self.state["input_feed"] = fn(self.state["input_feed"], 1)
+        if self._coverage and self.state["coverage"] is not None:
+            self.state["coverage"] = fn(self.state["coverage"], 1)
+
+    def detach_state(self):
+        self.state["hidden"] = tuple(h.detach() for h in self.state["hidden"])
+        self.state["input_feed"] = self.state["input_feed"].detach()
+        if self._coverage and self.state["coverage"] is not None:
+            self.state["coverage"] = self.state["coverage"].detach()
+
+    def forward(self, tgt, enc_out, src_len=None, step=None, **kwargs):
+        """
+        Args:
+            tgt (LongTensor): sequences of padded tokens
+                 ``(batch, tgt_len, nfeats)``.
+            enc_out (FloatTensor): vectors from the encoder
+                 ``(batch, src_len, hidden)``.
+            src_len (LongTensor): the padded source lengths
+                ``(batch,)``.
+
+        Returns:
+            (FloatTensor, dict[str, FloatTensor]):
+
+            * dec_outs: output from the decoder (after attn)
+              ``(batch, tgt_len, hidden)``.
+            * attns: distribution over src at each tgt
+              ``(batch, tgt_len, src_len)``.
+        """
+        dec_state, dec_outs, attns = self._run_forward_pass(
+            tgt, enc_out, src_len=src_len
+        )
+
+        # Update the state with the result.
+        if not isinstance(dec_state, tuple):
+            dec_state = (dec_state,)
+        self.state["hidden"] = dec_state
+
+        # Concatenates sequence of tensors along a new dimension.
+        # NOTE: v0.3 to 0.4: dec_outs / attns[*] may not be list
+        #       (in particular in case of SRU) it was not raising error in 0.3
+        #       since stack(Variable) was allowed.
+        #       In 0.4, SRU returns a tensor that shouldn't be stacke
+        if type(dec_outs) == list:
+            dec_outs = torch.stack(dec_outs, dim=1)
+            for k in attns:
+                if type(attns[k]) == list:
+                    attns[k] = torch.stack(attns[k])
+
+        self.state["input_feed"] = dec_outs[:, -1, :].unsqueeze(0)
+        self.state["coverage"] = None
+        if "coverage" in attns:
+            self.state["coverage"] = attns["coverage"][-1, :, :].unsqueeze(0)
+
+        return dec_outs, attns
+
+    def update_dropout(self, dropout, attention_dropout=None):
+        self.dropout.p = dropout
+        self.embeddings.update_dropout(dropout)
+
+
+class StdRNNDecoder(RNNDecoderBase):
+    """Standard fully batched RNN decoder with attention.
+
+    Faster implementation, uses CuDNN for implementation.
+    See :class:`~onmt.decoders.decoder.RNNDecoderBase` for options.
+
+
+    Based around the approach from
+    "Neural Machine Translation By Jointly Learning To Align and Translate"
+    :cite:`Bahdanau2015`
+
+
+    Implemented without input_feeding and currently with no `coverage_attn`
+    or `copy_attn` support.
+    """
+
+    def _run_forward_pass(self, tgt, enc_out, src_len=None):
+        """
+        Private helper for running the specific RNN forward pass.
+        Must be overriden by all subclasses.
+
+        Args:
+            tgt (LongTensor): a sequence of input tokens tensors
+                ``(batch, tgt_len, nfeats)``.
+            enc_out (FloatTensor): output(tensor sequence) from the
+                encoder RNN of size ``(batch, src_len, hidden_size)``.
+            src_len (LongTensor): the source enc_out lengths.
+
+        Returns:
+            (Tensor, List[FloatTensor], Dict[str, List[FloatTensor]):
+
+            * dec_state: final hidden state from the decoder.
+            * dec_outs: an array of output of every time
+              step from the decoder.
+            * attns: a dictionary of different
+              type of attention Tensor array of every time
+              step from the decoder.
+        """
+
+        assert self.copy_attn is None  # TODO, no support yet.
+        assert not self._coverage  # TODO, no support yet.
+
+        attns = {}
+        emb = self.embeddings(tgt)
+
+        if isinstance(self.rnn, nn.GRU):
+            rnn_out, dec_state = self.rnn(emb, self.state["hidden"][0])
+        else:
+            rnn_out, dec_state = self.rnn(emb, self.state["hidden"])
+
+        tgt_batch, tgt_len, _ = tgt.size()
+
+        # Calculate the attention.
+        if not self.attentional:
+            dec_outs = rnn_out
+        else:
+            dec_outs, p_attn = self.attn(rnn_out, enc_out, src_len=src_len)
+            attns["std"] = p_attn
+
+        # Calculate the context gate.
+        if self.context_gate is not None:
+            dec_outs = self.context_gate(
+                emb.view(-1, emb.size(2)),
+                rnn_out.view(-1, rnn_out.size(2)),
+                dec_outs.view(-1, dec_outs.size(2)),
+            )
+            dec_outs = dec_outs.view(tgt_batch, tgt_len, self.hidden_size)
+
+        dec_outs = self.dropout(dec_outs)
+
+        return dec_state, dec_outs, attns
+
+    def _build_rnn(self, rnn_type, **kwargs):
+        rnn, _ = rnn_factory(rnn_type, **kwargs)
+        return rnn
+
+    @property
+    def _input_size(self):
+        return self.embeddings.embedding_size
+
+
+class InputFeedRNNDecoder(RNNDecoderBase):
+    """Input feeding based decoder.
+
+    See :class:`~onmt.decoders.decoder.RNNDecoderBase` for options.
+
+    Based around the input feeding approach from
+    "Effective Approaches to Attention-based Neural Machine Translation"
+    :cite:`Luong2015`
+
+    """
+
+    def _run_forward_pass(self, tgt, enc_out, src_len=None):
+        """
+        See StdRNNDecoder._run_forward_pass() for description
+        of arguments and return values.
+        """
+        # Additional args check.
+        input_feed = self.state["input_feed"].squeeze(0)
+
+        dec_outs = []
+        attns = {}
+        if self.attn is not None:
+            attns["std"] = []
+        if self.copy_attn is not None or self._reuse_copy_attn:
+            attns["copy"] = []
+        if self._coverage:
+            attns["coverage"] = []
+
+        emb = self.embeddings(tgt)
+        assert emb.dim() == 3  # batch x len x embedding_dim
+
+        dec_state = self.state["hidden"]
+
+        coverage = (
+            self.state["coverage"].squeeze(0)
+            if self.state["coverage"] is not None
+            else None
+        )
+
+        # Input feed concatenates hidden state with
+        # input at every time step.
+        for emb_t in emb.split(1, dim=1):
+            dec_in = torch.cat([emb_t.squeeze(1), input_feed], 1)
+            rnn_out, dec_state = self.rnn(dec_in, dec_state)
+            if self.attentional:
+                dec_out, p_attn = self.attn(rnn_out, enc_out, src_len=src_len)
+                attns["std"].append(p_attn)
+            else:
+                dec_out = rnn_out
+            if self.context_gate is not None:
+                # TODO: context gate should be employed
+                # instead of second RNN transform.
+                dec_out = self.context_gate(dec_in, rnn_out, dec_out)
+            dec_out = self.dropout(dec_out)
+            input_feed = dec_out
+
+            dec_outs += [dec_out]
+
+            # Update the coverage attention.
+            # attns["coverage"] is actually c^(t+1) of See et al(2017)
+            # 1-index shifted
+            if self._coverage:
+                coverage = p_attn if coverage is None else p_attn + coverage
+                attns["coverage"] += [coverage]
+
+            if self.copy_attn is not None:
+                _, copy_attn = self.copy_attn(dec_out, enc_out)
+                attns["copy"] += [copy_attn]
+            elif self._reuse_copy_attn:
+                attns["copy"] = attns["std"]
+
+        return dec_state, dec_outs, attns
+
+    def _build_rnn(self, rnn_type, input_size, hidden_size, num_layers, dropout):
+        assert rnn_type != "SRU", (
+            "SRU doesn't support input feed! " "Please set -input_feed 0!"
+        )
+        stacked_cell = StackedLSTM if rnn_type == "LSTM" else StackedGRU
+        return stacked_cell(num_layers, input_size, hidden_size, dropout)
+
+    @property
+    def _input_size(self):
+        """Using input feed by concatenating input with attention vectors."""
+        return self.embeddings.embedding_size + self.hidden_size
+
+    def update_dropout(self, dropout, attention_dropout=None):
+        self.dropout.p = dropout
+        self.rnn.dropout.p = dropout
+        self.embeddings.update_dropout(dropout)

onmt/decoders/ensemble.py

@@ -0,0 +1,150 @@
+"""Ensemble decoding.
+
+Decodes using multiple models simultaneously,
+combining their prediction distributions by averaging.
+All models in the ensemble must share a target vocabulary.
+"""
+
+import torch
+import torch.nn as nn
+
+from onmt.encoders.encoder import EncoderBase
+from onmt.decoders.decoder import DecoderBase
+from onmt.models import NMTModel
+import onmt.model_builder
+
+
+class EnsembleDecoderOutput(object):
+    """Wrapper around multiple decoder final hidden states."""
+
+    def __init__(self, model_dec_outs):
+        self.model_dec_outs = tuple(model_dec_outs)
+
+    def squeeze(self, dim=None):
+        """Delegate squeeze to avoid modifying
+        :func:`onmt.translate.translator.Translator.translate_batch()`
+        """
+        return EnsembleDecoderOutput([x.squeeze(dim) for x in self.model_dec_outs])
+
+    def __getitem__(self, index):
+        return self.model_dec_outs[index]
+
+
+class EnsembleEncoder(EncoderBase):
+    """Dummy Encoder that delegates to individual real Encoders."""
+
+    def __init__(self, model_encoders):
+        super(EnsembleEncoder, self).__init__()
+        self.model_encoders = nn.ModuleList(model_encoders)
+
+    def forward(self, src, src_len=None):
+        enc_out, enc_final_hs, _ = zip(
+            *[model_encoder(src, src_len) for model_encoder in self.model_encoders]
+        )
+        return enc_out, enc_final_hs, src_len
+
+
+class EnsembleDecoder(DecoderBase):
+    """Dummy Decoder that delegates to individual real Decoders."""
+
+    def __init__(self, model_decoders):
+        model_decoders = nn.ModuleList(model_decoders)
+        attentional = any([dec.attentional for dec in model_decoders])
+        super(EnsembleDecoder, self).__init__(attentional)
+        self.model_decoders = model_decoders
+
+    def forward(self, tgt, enc_out, src_len=None, step=None, **kwargs):
+        """See :func:`onmt.decoders.decoder.DecoderBase.forward()`."""
+        # src_len is a single tensor shared between all models.
+        # This assumption will not hold if Translator is modified
+        # to calculate src_len as something other than the length
+        # of the input.
+        dec_outs, attns = zip(
+            *[
+                model_decoder(tgt, enc_out[i], src_len=src_len, step=step, **kwargs)
+                for i, model_decoder in enumerate(self.model_decoders)
+            ]
+        )
+        mean_attns = self.combine_attns(attns)
+        return EnsembleDecoderOutput(dec_outs), mean_attns
+
+    def combine_attns(self, attns):
+        result = {}
+        for key in attns[0].keys():
+            result[key] = torch.stack(
+                [attn[key] for attn in attns if attn[key] is not None]
+            ).mean(0)
+        return result
+
+    def init_state(self, src, enc_out, enc_hidden):
+        """See :obj:`RNNDecoderBase.init_state()`"""
+        for i, model_decoder in enumerate(self.model_decoders):
+            model_decoder.init_state(src, enc_out[i], enc_hidden[i])
+
+    def map_state(self, fn):
+        for model_decoder in self.model_decoders:
+            model_decoder.map_state(fn)
+
+
+class EnsembleGenerator(nn.Module):
+    """
+    Dummy Generator that delegates to individual real Generators,
+    and then averages the resulting target distributions.
+    """
+
+    def __init__(self, model_generators, raw_probs=False):
+        super(EnsembleGenerator, self).__init__()
+        self.model_generators = nn.ModuleList(model_generators)
+        self._raw_probs = raw_probs
+
+    def forward(self, hidden, attn=None, src_map=None):
+        """
+        Compute a distribution over the target dictionary
+        by averaging distributions from models in the ensemble.
+        All models in the ensemble must share a target vocabulary.
+        """
+        distributions = torch.stack(
+            [
+                mg(h) if attn is None else mg(h, attn, src_map)
+                for h, mg in zip(hidden, self.model_generators)
+            ]
+        )
+        if self._raw_probs:
+            return torch.log(torch.exp(distributions).mean(0))
+        else:
+            return distributions.mean(0)
+
+
+class EnsembleModel(NMTModel):
+    """Dummy NMTModel wrapping individual real NMTModels."""
+
+    def __init__(self, models, raw_probs=False):
+        encoder = EnsembleEncoder(model.encoder for model in models)
+        decoder = EnsembleDecoder(model.decoder for model in models)
+        super(EnsembleModel, self).__init__(encoder, decoder)
+        self.generator = EnsembleGenerator(
+            [model.generator for model in models], raw_probs
+        )
+        self.models = nn.ModuleList(models)
+
+
+def load_test_model(opt, device_id=0):
+    """Read in multiple models for ensemble."""
+    shared_vocabs = None
+    shared_model_opt = None
+    models = []
+    for model_path in opt.models:
+        vocabs, model, model_opt = onmt.model_builder.load_test_model(
+            opt, device_id, model_path=model_path
+        )
+        if shared_vocabs is None:
+            shared_vocabs = vocabs
+        else:
+            assert (
+                shared_vocabs["src"].tokens_to_ids == vocabs["src"].tokens_to_ids
+            ), "Ensemble models must use the same vocabs "
+        models.append(model)
+        if shared_model_opt is None:
+            shared_model_opt = model_opt
+    ensemble_model = EnsembleModel(models, opt.avg_raw_probs)
+    return shared_vocabs, ensemble_model, shared_model_opt

onmt/decoders/transformer.py

@@ -0,0 +1,835 @@
+"""
+Implementation of "Attention is All You Need" and of
+subsequent transformer based architectures
+"""
+
+import torch
+import torch.nn as nn
+from onmt.decoders.decoder import DecoderBase
+from onmt.modules import MultiHeadedAttention, AverageAttention
+from onmt.modules.position_ffn import PositionwiseFeedForward
+from onmt.modules.position_ffn import ActivationFunction
+from onmt.utils.misc import sequence_mask
+from onmt.modules.rmsnorm import RMSNorm
+
+
+class TransformerDecoderLayerBase(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        heads,
+        d_ff,
+        dropout,
+        attention_dropout,
+        self_attn_type="scaled-dot",
+        max_relative_positions=0,
+        relative_positions_buckets=0,
+        aan_useffn=False,
+        full_context_alignment=False,
+        alignment_heads=0,
+        pos_ffn_activation_fn=ActivationFunction.relu,
+        add_qkvbias=False,
+        num_kv=0,
+        add_ffnbias=True,
+        parallel_residual=False,
+        shared_layer_norm=False,
+        layer_norm="standard",
+        norm_eps=1e-6,
+        use_ckpting=[],
+        parallel_gpu=1,
+    ):
+        """
+        Args:
+            d_model (int): the dimension of keys/values/queries in
+                :class:`MultiHeadedAttention`, also the input size of
+                the first-layer of the :class:`PositionwiseFeedForward`.
+            heads (int): the number of heads for MultiHeadedAttention.
+            d_ff (int): the second-layer of the
+                :class:`PositionwiseFeedForward`.
+            dropout (float): dropout in residual, self-attn(dot) and
+                feed-forward
+            attention_dropout (float): dropout in context_attn  (and
+                self-attn(avg))
+            self_attn_type (string): type of self-attention scaled-dot,
+                average
+            max_relative_positions (int):
+                Max distance between inputs in relative positions
+                representations
+            aan_useffn (bool): Turn on the FFN layer in the AAN decoder
+            full_context_alignment (bool):
+                whether enable an extra full context decoder forward for
+                alignment
+            alignment_heads (int):
+                N. of cross attention heads to use for alignment guiding
+            pos_ffn_activation_fn (ActivationFunction):
+                activation function choice for PositionwiseFeedForward layer
+            add_qkvbias (bool): whether to add bias to the Key/Value nn.Linear
+            layer_norm (string): type of layer normalization standard/rms
+            norm_eps (float): layer norm epsilon
+
+        """
+        super(TransformerDecoderLayerBase, self).__init__()
+
+        self.self_attn_type = self_attn_type
+        if self_attn_type == "scaled-dot":
+            self.self_attn = MultiHeadedAttention(
+                heads,
+                d_model,
+                dropout=attention_dropout,
+                max_relative_positions=max_relative_positions,
+                relative_positions_buckets=relative_positions_buckets,
+                attn_type="self",
+                add_qkvbias=add_qkvbias,
+                num_kv=num_kv,
+                use_ckpting=use_ckpting,
+                parallel_gpu=parallel_gpu,
+            )
+        elif self_attn_type == "average":
+            self.self_attn = AverageAttention(
+                d_model, dropout=attention_dropout, aan_useffn=aan_useffn
+            )
+
+        self.feed_forward = PositionwiseFeedForward(
+            d_model,
+            d_ff,
+            dropout,
+            pos_ffn_activation_fn,
+            add_ffnbias,
+            parallel_residual,
+            layer_norm,
+            norm_eps,
+            use_ckpting=use_ckpting,
+            parallel_gpu=parallel_gpu,
+        )
+        self.parallel_residual = parallel_residual
+        self.shared_layer_norm = shared_layer_norm
+        if layer_norm == "standard":
+            self.layer_norm_1 = nn.LayerNorm(d_model, eps=norm_eps)
+            if parallel_residual and not shared_layer_norm:
+                self.layer_norm_res = nn.LayerNorm(d_model, eps=norm_eps)
+        elif layer_norm == "rms":
+            self.layer_norm_1 = RMSNorm(d_model, eps=norm_eps)
+            if parallel_residual and not shared_layer_norm:
+                self.layer_norm_res = RMSNorm(d_model, eps=norm_eps)
+        else:
+            raise ValueError(f"{layer_norm} layer norm type is not supported")
+
+        self.dropout = nn.Dropout(dropout)
+        self.full_context_alignment = full_context_alignment
+        self.alignment_heads = alignment_heads
+
+    def forward(self, *args, **kwargs):
+        """Extend `_forward` for (possibly) multiple decoder pass:
+        Always a default (future masked) decoder forward pass,
+        Possibly a second future aware decoder pass for joint learn
+        full context alignement, :cite:`garg2019jointly`.
+
+        Args:
+            * All arguments of _forward, of which
+            with_align (bool): needed to compute attn_align
+            return_attn (bool): to force MHA to return attns
+
+        Returns:
+            (FloatTensor, FloatTensor, FloatTensor or None):
+
+            * layer_out ``(batch_size, T, model_dim)``
+            * top_attn ``(batch_size, T, src_len)``
+            * attn_align ``(batch_size, T, src_len)`` or None
+        """
+        with_align = kwargs.pop("with_align", False)
+        layer_out, attns = self._forward(*args, **kwargs)
+        top_attn = None if attns is None else attns[:, 0, :, :].contiguous()
+        attn_align = None
+        if with_align:
+            if self.full_context_alignment:
+                # return _, (B, Q_len, K_len)
+                _, attns = self._forward(*args, **kwargs, future=True)
+
+            if self.alignment_heads > 0:
+                attns = attns[:, : self.alignment_heads, :, :].contiguous()
+            # layer average attention across heads, get ``(B, Q, K)``
+            # Case 1: no full_context, no align heads -> layer avg baseline
+            # Case 2: no full_context, 1 align heads -> guided align
+            # Case 3: full_context, 1 align heads -> full cte guided align
+            attn_align = attns.mean(dim=1)
+        return layer_out, top_attn, attn_align
+
+    def update_dropout(self, dropout, attention_dropout):
+        self.self_attn.update_dropout(attention_dropout)
+        self.feed_forward.update_dropout(dropout)
+        self.dropout.p = dropout
+
+    def _forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def _compute_dec_mask(self, tgt_pad_mask, future):
+        tgt_len = tgt_pad_mask.size(-1)
+        if not future:  # apply future_mask, result mask in (B, T, T)
+            future_mask = torch.ones(
+                [tgt_len, tgt_len],
+                device=tgt_pad_mask.device,
+                dtype=torch.uint8,
+            )
+            future_mask = future_mask.triu_(1).view(1, tgt_len, tgt_len)
+            # BoolTensor was introduced in pytorch 1.2
+            try:
+                future_mask = future_mask.bool()
+            except AttributeError:
+                pass
+            dec_mask = torch.gt(tgt_pad_mask + future_mask, 0)
+        else:  # only mask padding, result mask in (B, 1, T)
+            dec_mask = tgt_pad_mask
+        return dec_mask
+
+    def _forward_self_attn(self, norm_layer_in, dec_mask, step, return_attn=False):
+        if self.self_attn_type == "scaled-dot":
+            return self.self_attn(
+                norm_layer_in,
+                norm_layer_in,
+                norm_layer_in,
+                mask=dec_mask,
+                step=step,
+                return_attn=return_attn,
+            )
+        elif self.self_attn_type == "average":
+            return self.self_attn(norm_layer_in, mask=dec_mask, step=step)
+        else:
+            raise ValueError(f"self attention {type(self.self_attn)} not supported")
+
+
+class TransformerDecoderLayer(TransformerDecoderLayerBase):
+    """Transformer Decoder layer block in Pre-Norm style.
+    Pre-Norm style is an improvement w.r.t. Original paper's Post-Norm style,
+    providing better converge speed and performance. This is also the actual
+    implementation in tensor2tensor and also avalable in fairseq.
+    See https://tunz.kr/post/4 and :cite:`DeeperTransformer`.
+
+    """
+
+    def __init__(
+        self,
+        d_model,
+        heads,
+        d_ff,
+        dropout,
+        attention_dropout,
+        self_attn_type="scaled-dot",
+        max_relative_positions=0,
+        relative_positions_buckets=0,
+        aan_useffn=False,
+        full_context_alignment=False,
+        alignment_heads=0,
+        pos_ffn_activation_fn=ActivationFunction.relu,
+        add_qkvbias=False,
+        num_kv=0,
+        add_ffnbias=True,
+        parallel_residual=False,
+        shared_layer_norm=False,
+        layer_norm="standard",
+        norm_eps=1e-6,
+        use_ckpting=[],
+        parallel_gpu=1,
+    ):
+        """
+        Args:
+            See TransformerDecoderLayerBase
+        """
+        super(TransformerDecoderLayer, self).__init__(
+            d_model,
+            heads,
+            d_ff,
+            dropout,
+            attention_dropout,
+            self_attn_type,
+            max_relative_positions,
+            relative_positions_buckets,
+            aan_useffn,
+            full_context_alignment,
+            alignment_heads,
+            pos_ffn_activation_fn=pos_ffn_activation_fn,
+            add_qkvbias=add_qkvbias,
+            num_kv=num_kv,
+            add_ffnbias=add_ffnbias,
+            parallel_residual=parallel_residual,
+            shared_layer_norm=shared_layer_norm,
+            layer_norm=layer_norm,
+            norm_eps=norm_eps,
+            use_ckpting=use_ckpting,
+            parallel_gpu=parallel_gpu,
+        )
+        self.context_attn = MultiHeadedAttention(
+            heads,
+            d_model,
+            dropout=attention_dropout,
+            attn_type="context",
+            add_qkvbias=add_qkvbias,
+            num_kv=num_kv,
+            use_ckpting=use_ckpting,
+            parallel_gpu=parallel_gpu,
+        )
+        if layer_norm == "standard":
+            self.layer_norm_2 = nn.LayerNorm(d_model, eps=norm_eps)
+        elif layer_norm == "rms":
+            self.layer_norm_2 = RMSNorm(d_model, eps=norm_eps)
+        else:
+            raise ValueError(f"{layer_norm} layer norm type is not supported")
+
+    def update_dropout(self, dropout, attention_dropout):
+        super(TransformerDecoderLayer, self).update_dropout(dropout, attention_dropout)
+        self.context_attn.update_dropout(attention_dropout)
+
+    def _forward(
+        self,
+        layer_in,
+        enc_out,
+        src_pad_mask,
+        tgt_pad_mask,
+        step=None,
+        future=False,
+        return_attn=False,
+    ):
+        """A naive forward pass for transformer decoder.
+
+        # T: could be 1 in the case of stepwise decoding or tgt_len
+
+        Args:
+            layer_in (FloatTensor): ``(batch_size, T, model_dim)``
+            enc_out (FloatTensor): ``(batch_size, src_len, model_dim)``
+            src_pad_mask (bool): ``(batch_size, 1, src_len)``
+            tgt_pad_mask (bool): ``(batch_size, 1, T)``
+            step (int or None): stepwise decoding counter
+            future (bool): If set True, do not apply future_mask.
+            return_attn (bool) : if set True requires attns output
+
+        Returns:
+            (FloatTensor, FloatTensor):
+
+            * layer_out ``(batch_size, T, model_dim)``
+            * attns ``(batch_size, head, T, src_len)``
+
+        """
+        dec_mask = None
+        src_pad_mask = src_pad_mask.unsqueeze(1)  # [B,1,1,slen]
+
+        if layer_in.size(1) > 1:
+            # masking is necessary when sequence length is greater than one
+            dec_mask = self._compute_dec_mask(tgt_pad_mask, future)
+            dec_mask = dec_mask.unsqueeze(1)
+            dec_mask = dec_mask.expand(-1, -1, dec_mask.size(3), -1)
+            src_pad_mask = src_pad_mask.expand(-1, -1, dec_mask.size(3), -1)
+            # mask now are (batch x 1 x tlen x s or t len)
+            # 1 = heads to be expanded in MHA
+
+        norm_layer_in = self.layer_norm_1(layer_in)
+
+        self_attn, _ = self._forward_self_attn(norm_layer_in, dec_mask, step)
+
+        if self.parallel_residual:
+            ctx_attn, attns = self.context_attn(
+                enc_out,
+                enc_out,
+                norm_layer_in,
+                mask=src_pad_mask,
+                return_attn=return_attn,
+            )
+            # feed_forward applies residual, so we remove and apply residual with un-normed
+            layer_out = (
+                self.feed_forward(norm_layer_in)
+                - norm_layer_in
+                + layer_in
+                + self.dropout(self_attn)
+                + ctx_attn
+            )
+        else:
+            query = self.dropout(self_attn) + layer_in
+            norm_query = self.layer_norm_2(query)
+            ctx_attn, attns = self.context_attn(
+                enc_out, enc_out, norm_query, mask=src_pad_mask, return_attn=return_attn
+            )
+            layer_out = self.feed_forward(self.dropout(ctx_attn) + query)
+
+        return layer_out, attns
+
+
+class TransformerDecoderBase(DecoderBase):
+    def __init__(
+        self, d_model, copy_attn, embeddings, alignment_layer, layer_norm, norm_eps
+    ):
+        super(TransformerDecoderBase, self).__init__()
+
+        self.embeddings = embeddings
+
+        # Decoder State
+        self.state = {}
+
+        # previously, there was a GlobalAttention module here for copy
+        # attention. But it was never actually used -- the "copy" attention
+        # just reuses the context attention.
+        self._copy = copy_attn
+        if layer_norm == "standard":
+            self.layer_norm = nn.LayerNorm(d_model, eps=norm_eps)
+        elif layer_norm == "rms":
+            self.layer_norm = RMSNorm(d_model, eps=norm_eps)
+        else:
+            raise ValueError(f"{layer_norm} layer norm type is not supported")
+
+        self.alignment_layer = alignment_layer
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor."""
+        return cls(
+            opt.dec_layers,
+            opt.dec_hid_size,
+            opt.heads,
+            opt.transformer_ff,
+            opt.copy_attn,
+            opt.self_attn_type,
+            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
+            opt.attention_dropout[0]
+            if type(opt.attention_dropout) is list
+            else opt.attention_dropout,
+            embeddings,
+            opt.max_relative_positions,
+            opt.relative_positions_buckets,
+            opt.aan_useffn,
+            opt.full_context_alignment,
+            opt.alignment_layer,
+            alignment_heads=opt.alignment_heads,
+            pos_ffn_activation_fn=opt.pos_ffn_activation_fn,
+            add_qkvbias=opt.add_qkvbias,
+            num_kv=opt.num_kv,
+            add_ffnbias=opt.add_ffnbias,
+            parallel_residual=opt.parallel_residual,
+            shared_layer_norm=opt.shared_layer_norm,
+            layer_norm=opt.layer_norm,
+            norm_eps=opt.norm_eps,
+            use_ckpting=opt.use_ckpting,
+            parallel_gpu=opt.world_size
+            if opt.parallel_mode == "tensor_parallel"
+            else 1,
+        )
+
+    def init_state(self, src, enc_out, enc_final_hs):
+        """Initialize decoder state."""
+        self.state["src"] = src
+
+    def map_state(self, fn):
+        if self.state["src"] is not None:
+            self.state["src"] = fn(self.state["src"], 0)
+        for layer in self.transformer_layers:
+            if hasattr(layer, "context_attn"):
+                if layer.context_attn.layer_cache[1]["keys"].numel() != 0:
+                    x = fn(layer.context_attn.layer_cache[1]["keys"], 0)
+                    y = fn(layer.context_attn.layer_cache[1]["values"], 0)
+                    layer.context_attn.layer_cache = True, {"keys": x, "values": y}
+            if isinstance(layer.self_attn, AverageAttention):
+                if layer.self_attn.layer_cache[1]["prev_g"].numel() != 0:
+                    x = fn(layer.self_attn.layer_cache[1]["prev_g"], 0)
+                    layer.self_attn.layer_cache = True, {"prev_g": x}
+            else:
+                if layer.self_attn.layer_cache[1]["keys"].numel() != 0:
+                    x = fn(layer.self_attn.layer_cache[1]["keys"], 0)
+                    y = fn(layer.self_attn.layer_cache[1]["values"], 0)
+                    layer.self_attn.layer_cache = True, {"keys": x, "values": y}
+
+    def detach_state(self):
+        raise NotImplementedError
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def update_dropout(self, dropout, attention_dropout):
+        self.embeddings.update_dropout(dropout)
+        for layer in self.transformer_layers:
+            layer.update_dropout(dropout, attention_dropout)
+
+
+class TransformerDecoder(TransformerDecoderBase):
+    """The Transformer decoder from "Attention is All You Need".
+    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`
+
+    Args:
+        num_layers (int): number of decoder layers.
+        d_model (int): size of the model
+        heads (int): number of heads
+        d_ff (int): size of the inner FF layer
+        copy_attn (bool): if using a separate copy attention
+        self_attn_type (str): type of self-attention scaled-dot, average
+        dropout (float): dropout in residual, self-attn(dot) and feed-forward
+        attention_dropout (float): dropout in context_attn (and self-attn(avg))
+        embeddings (onmt.modules.Embeddings):
+            embeddings to use, should have positional encodings
+        max_relative_positions (int):
+            Max distance between inputs in relative positions representations
+        relative_positions_buckets (int):
+            Number of buckets when using relative position bias
+        aan_useffn (bool): Turn on the FFN layer in the AAN decoder
+        full_context_alignment (bool):
+            whether enable an extra full context decoder forward for alignment
+        alignment_layer (int): N° Layer to supervise with for alignment guiding
+        alignment_heads (int):
+            N. of cross attention heads to use for alignment guiding
+        add_qkvbias (bool): whether to add bias to the Key/Value nn.Linear
+        layer_norm (string): type of layer normalization standard/rms
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        d_model,
+        heads,
+        d_ff,
+        copy_attn,
+        self_attn_type,
+        dropout,
+        attention_dropout,
+        embeddings,
+        max_relative_positions,
+        relative_positions_buckets,
+        aan_useffn,
+        full_context_alignment,
+        alignment_layer,
+        alignment_heads,
+        pos_ffn_activation_fn=ActivationFunction.relu,
+        add_qkvbias=False,
+        num_kv=0,
+        add_ffnbias=True,
+        parallel_residual=False,
+        shared_layer_norm=False,
+        layer_norm="standard",
+        norm_eps=1e-6,
+        use_ckpting=[],
+        parallel_gpu=1,
+    ):
+        super(TransformerDecoder, self).__init__(
+            d_model, copy_attn, embeddings, alignment_layer, layer_norm, norm_eps
+        )
+
+        self.transformer_layers = nn.ModuleList(
+            [
+                TransformerDecoderLayer(
+                    d_model,
+                    heads,
+                    d_ff,
+                    dropout,
+                    attention_dropout,
+                    self_attn_type=self_attn_type,
+                    max_relative_positions=max_relative_positions,
+                    relative_positions_buckets=relative_positions_buckets,
+                    aan_useffn=aan_useffn,
+                    full_context_alignment=full_context_alignment,
+                    alignment_heads=alignment_heads,
+                    pos_ffn_activation_fn=pos_ffn_activation_fn,
+                    add_qkvbias=add_qkvbias,
+                    num_kv=num_kv,
+                    add_ffnbias=add_ffnbias,
+                    parallel_residual=parallel_residual,
+                    shared_layer_norm=shared_layer_norm,
+                    layer_norm=layer_norm,
+                    norm_eps=norm_eps,
+                    use_ckpting=use_ckpting,
+                    parallel_gpu=parallel_gpu,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+    def detach_state(self):
+        self.state["src"] = self.state["src"].detach()
+
+    def forward(self, tgt, enc_out=None, step=None, **kwargs):
+        """
+        Decode, possibly stepwise.
+        when training step is always None, when decoding, step increases
+        tgt (Tensor): batch x tlen x feats
+        enc_out (Tensor): encoder output (batch x slen x model_dim)
+        """
+        if enc_out is None:
+            enc_out = self.embeddings(tgt)
+        if step == 0:
+            self._init_cache(enc_out)
+        elif step is None:
+            for layer in self.transformer_layers:
+                if isinstance(layer.self_attn, AverageAttention):
+                    layer.self_attn.layer_cache = False, {"prev_g": torch.tensor([])}
+                else:
+                    layer.self_attn.layer_cache = (
+                        False,
+                        {"keys": torch.tensor([]), "values": torch.tensor([])},
+                    )
+                layer.context_attn.layer_cache = (
+                    False,
+                    {"keys": torch.tensor([]), "values": torch.tensor([])},
+                )
+
+        emb = self.embeddings(tgt, step=step)
+        dec_out = emb
+        assert emb.dim() == 3  # len x batch x embedding_dim
+
+        pad_idx = self.embeddings.word_padding_idx
+        src_lens = kwargs["src_len"]
+        src_max_len = self.state["src"].shape[1]
+        src_pad_mask = ~sequence_mask(src_lens, src_max_len)  # [B x slen]
+        src_pad_mask = src_pad_mask.unsqueeze(1)  # [B x 1 x slen]
+        tgt_pad_mask = tgt[:, :, 0].eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]
+
+        with_align = kwargs.pop("with_align", False)
+        return_attn = with_align or self._copy
+        attn_aligns = []
+
+        for layer in self.transformer_layers:
+            dec_out, attn, attn_align = layer(
+                dec_out,
+                enc_out,
+                src_pad_mask,
+                tgt_pad_mask,
+                step=step,
+                with_align=with_align,
+                return_attn=return_attn,
+            )
+            if attn_align is not None:
+                attn_aligns.append(attn_align)
+
+        dec_out = self.layer_norm(dec_out)
+
+        attns = {"std": attn}
+        if self._copy:
+            attns["copy"] = attn
+        if with_align:
+            attns["align"] = attn_aligns[self.alignment_layer]  # `(B, Q, K)`
+            # attns["align"] = torch.stack(attn_aligns, 0).mean(0)  # All avg
+
+        # TODO change the way attns is returned dict => list or tuple (onnx)
+        return dec_out, attns
+
+    def _init_cache(self, enc_out):
+        batch_size = enc_out.size(0)
+        depth = enc_out.size(-1)
+
+        for layer in self.transformer_layers:
+            # first value set to True triggered by the beginning of decoding
+            # layer_cache becomes active in the MultiHeadedAttention fwd
+            layer.context_attn.layer_cache = (
+                True,
+                {
+                    "keys": torch.tensor([], device=enc_out.device),
+                    "values": torch.tensor([], device=enc_out.device),
+                },
+            )
+            if isinstance(layer.self_attn, AverageAttention):
+                layer.self_attn.layer_cache = True, {
+                    "prev_g": torch.zeros(
+                        (batch_size, 1, depth), device=enc_out.device
+                    ).to(enc_out.dtype)
+                }
+            else:
+                layer.self_attn.layer_cache = (
+                    True,
+                    {
+                        "keys": torch.tensor([], device=enc_out.device),
+                        "values": torch.tensor([], device=enc_out.device),
+                    },
+                )
+
+
+class TransformerLMDecoderLayer(TransformerDecoderLayerBase):
+    """Transformer Decoder only layer block in GPT style.
+    Args:
+         See TransformerDecoderLayerBase
+    """
+
+    def _forward(
+        self, layer_in, tgt_pad_mask, step=None, future=False, return_attn=False
+    ):
+        """A naive forward pass for transformer decoder.
+
+        # T: could be 1 in the case of stepwise decoding or tgt_len
+
+        Args:
+            layer_in (FloatTensor): ``(batch_size, T, model_dim)``
+            tgt_pad_mask (bool): ``(batch_size, 1, T)``
+            layer_cache (dict or None): cached layer info when stepwise decode
+            step (int or None): stepwise decoding counter
+            future (bool): If set True, do not apply future_mask.
+            return_attn (bool): If set True return attn
+
+        Returns:
+            (FloatTensor, FloatTensor):
+
+            * layer_out ``(batch_size, T, model_dim)``
+            * attns ``(batch_size, head, T, T)``
+
+        """
+        dec_mask = None
+
+        if layer_in.size(1) > 1:
+            # masking is necessary when sequence length is greater than one
+            dec_mask = self._compute_dec_mask(tgt_pad_mask, future)
+            dec_mask = dec_mask.unsqueeze(1)
+            dec_mask = dec_mask.expand(-1, -1, dec_mask.size(3), -1)
+            # mask now are (batch x 1 x tlen x tlen)
+            # 1 = heads to be expanded in MHA
+
+        norm_layer_in = self.layer_norm_1(layer_in)
+
+        attn_output, attns = self._forward_self_attn(
+            norm_layer_in, dec_mask, step, return_attn=return_attn
+        )
+
+        if self.parallel_residual:
+            # feed_forward applies residual, so we remove and apply residual with un-normed
+            if not self.shared_layer_norm:
+                norm_res_layer_in = self.layer_norm_res(layer_in)
+                ff_in = norm_res_layer_in
+            else:
+                ff_in = norm_layer_in
+            layer_out = (
+                self.feed_forward(ff_in) - ff_in + layer_in + self.dropout(attn_output)
+            )
+        else:
+            layer_out = self.dropout(attn_output) + layer_in
+            layer_out = self.feed_forward(layer_out)
+
+        return layer_out, attns
+
+
+class TransformerLMDecoder(TransformerDecoderBase):
+    """The Transformer decoder from GPT-2
+    Args:
+         num_layers (int): number of decoder layers.
+         d_model (int): size of the model
+         heads (int): number of heads
+         d_ff (int): size of the inner FF layer
+         copy_attn (bool): if using a separate copy attention
+         self_attn_type (str): type of self-attention scaled-dot, average
+         dropout (float): dropout in residual, self-attn(dot) and feed-forward
+         attention_dropout (float): dropout in context_attn (and self-attn(avg))
+         embeddings (onmt.modules.Embeddings):
+             embeddings to use, should have positional encodings
+         max_relative_positions (int):
+             Max distance between inputs in relative positions representations
+         relative_positions_buckets (int):
+             Number of buckets when using Relative positions bias
+         aan_useffn (bool): Turn on the FFN layer in the AAN decoder
+         add_qkvbias (bool): whether to add bias to the Key/Value nn.Linear
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        d_model,
+        heads,
+        d_ff,
+        copy_attn,
+        self_attn_type,
+        dropout,
+        attention_dropout,
+        embeddings,
+        max_relative_positions,
+        relative_positions_buckets,
+        aan_useffn,
+        full_context_alignment=None,
+        alignment_layer=None,
+        alignment_heads=None,
+        pos_ffn_activation_fn=ActivationFunction.relu,
+        add_qkvbias=False,
+        num_kv=0,
+        add_ffnbias=True,
+        parallel_residual=False,
+        shared_layer_norm=False,
+        layer_norm="standard",
+        norm_eps=1e-6,
+        use_ckpting=[],
+        parallel_gpu=1,
+    ):
+        super(TransformerLMDecoder, self).__init__(
+            d_model, copy_attn, embeddings, alignment_layer, layer_norm, norm_eps
+        )
+        self.transformer_layers = nn.ModuleList(
+            [
+                TransformerLMDecoderLayer(
+                    d_model,
+                    heads,
+                    d_ff,
+                    dropout,
+                    attention_dropout,
+                    self_attn_type=self_attn_type,
+                    max_relative_positions=max_relative_positions,
+                    relative_positions_buckets=relative_positions_buckets,
+                    aan_useffn=aan_useffn,
+                    full_context_alignment=None,
+                    alignment_heads=None,
+                    pos_ffn_activation_fn=pos_ffn_activation_fn,
+                    add_qkvbias=add_qkvbias,
+                    num_kv=num_kv,
+                    add_ffnbias=add_ffnbias,
+                    parallel_residual=parallel_residual,
+                    shared_layer_norm=shared_layer_norm,
+                    layer_norm=layer_norm,
+                    norm_eps=norm_eps,
+                    use_ckpting=use_ckpting,
+                    parallel_gpu=parallel_gpu,
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+    def init_state(self, src=None, enc_out=None, enc_final_hs=None):
+        super(TransformerLMDecoder, self).init_state(None, None, None)
+
+    def detach_state(self):
+        pass
+
+    def forward(self, tgt, enc_out=None, step=None, **kwargs):
+        """Decode, possibly stepwise."""
+        if step == 0:
+            self._init_cache(tgt)
+        elif step is None:
+            for layer in self.transformer_layers:
+                layer.self_attn.layer_cache = (
+                    False,
+                    {"keys": torch.tensor([]), "values": torch.tensor([])},
+                )
+
+        dec_out = self.embeddings(tgt, step=step)
+
+        assert dec_out.dim() == 3  # batch x len x embedding_dim
+
+        pad_idx = self.embeddings.word_padding_idx
+        tgt_pad_mask = tgt[:, :, 0].eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]
+
+        with_align = kwargs.pop("with_align", False)
+        return_attn = with_align or self._copy
+        assert not with_align, "TransformerLMDecoder does not support align"
+
+        for layer in self.transformer_layers:
+            dec_out, attn, _ = layer(
+                dec_out,
+                tgt_pad_mask,
+                step=step,
+                with_align=with_align,
+                return_attn=return_attn,
+            )
+
+        dec_out = self.layer_norm(dec_out)
+
+        attns = {"std": attn}
+        if self._copy:
+            attns["copy"] = attn
+
+        # TODO change the way attns is returned dict => list or tuple (onnx)
+        return dec_out, attns
+
+    def _init_cache(self, tgt=None):
+        for layer in self.transformer_layers:
+            if isinstance(layer.self_attn, AverageAttention):
+                raise NotImplementedError
+            else:
+                layer.self_attn.layer_cache = (
+                    True,
+                    {
+                        "keys": torch.tensor([], device=tgt.device),
+                        "values": torch.tensor([], device=tgt.device),
+                    },
+                )

onmt/encoders/__init__.py

@@ -0,0 +1,67 @@
+"""Module defining encoders."""
+import os
+import importlib
+from onmt.encoders.encoder import EncoderBase
+from onmt.encoders.transformer import TransformerEncoder
+from onmt.encoders.ggnn_encoder import GGNNEncoder
+from onmt.encoders.rnn_encoder import RNNEncoder
+from onmt.encoders.cnn_encoder import CNNEncoder
+from onmt.encoders.mean_encoder import MeanEncoder
+
+
+str2enc = {
+    "ggnn": GGNNEncoder,
+    "rnn": RNNEncoder,
+    "brnn": RNNEncoder,
+    "cnn": CNNEncoder,
+    "transformer": TransformerEncoder,
+    "mean": MeanEncoder,
+}
+
+__all__ = [
+    "EncoderBase",
+    "TransformerEncoder",
+    "GGNNEncoder",
+    "RNNEncoder",
+    "CNNEncoder",
+    "MeanEncoder",
+    "str2enc",
+]
+
+
+def get_encoders_cls(encoder_names):
+    """Return valid encoder class indicated in `encoder_names`."""
+    encoders_cls = {}
+    for name in encoder_names:
+        if name not in str2enc:
+            raise ValueError("%s encoder not supported!" % name)
+        encoders_cls[name] = str2enc[name]
+    return encoders_cls
+
+
+def register_encoder(name):
+    """Encoder register that can be used to add new encoder class."""
+
+    def register_encoder_cls(cls):
+        if name in str2enc:
+            raise ValueError("Cannot register duplicate encoder ({})".format(name))
+        if not issubclass(cls, EncoderBase):
+            raise ValueError(f"encoder ({name}: {cls.__name_}) must extend EncoderBase")
+        str2enc[name] = cls
+        __all__.append(cls.__name__)  # added to be complete
+        return cls
+
+    return register_encoder_cls
+
+
+# Auto import python files in this directory
+encoder_dir = os.path.dirname(__file__)
+for file in os.listdir(encoder_dir):
+    path = os.path.join(encoder_dir, file)
+    if (
+        not file.startswith("_")
+        and not file.startswith(".")
+        and (file.endswith(".py") or os.path.isdir(path))
+    ):
+        file_name = file[: file.find(".py")] if file.endswith(".py") else file
+        module = importlib.import_module("onmt.encoders." + file_name)