transformer_helpers.py

import torch.nn
from models import TransformerEncDecModel, TransformerDecModel
from interfaces import TransformerEncDecInterface, TransformerDecOnlyInterface


def create_model(
    in_vocab_size,
    out_vocab_size,
    vec_dim,
    n_heads,
    encoder_n_layers,
    decoder_n_layers,
    is_null_encoder=False,
    mode="enc_dec",
) -> torch.nn.Module:
    args = dict(embedding_init="xavier", scale_mode="opennmt", mode=mode)
    if is_null_encoder:
        return TransformerDecModel(
            in_vocab_size,
            out_vocab_size,
            vec_dim,
            n_heads,
            num_encoder_layers=encoder_n_layers,
            num_decoder_layers=decoder_n_layers,
            tied_embedding=True,
            **args
        )
    else:
        return TransformerEncDecModel(
            in_vocab_size,
            out_vocab_size,
            vec_dim,
            n_heads,
            num_encoder_layers=encoder_n_layers,
            num_decoder_layers=decoder_n_layers,
            tied_embedding=True,
            **args
        )


def create_model_interface(model, label_smoothing=0.0, is_null_encoder=False):
    if is_null_encoder:
        return TransformerDecOnlyInterface(model, label_smoothing=label_smoothing)
    else:
        return TransformerEncDecInterface(model, label_smoothing=label_smoothing)