<gh_stars>0
from fairseq.models import register_model_architecture
from fairseq.models.bart import mbart_base_architecture

@register_model_architecture('bart', 'mbart_small')
def mbart_small(args):
    args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 768 // 4)
    args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4 * 768 // 4)
    args.encoder_layers = getattr(args, "encoder_layers", 3)
    args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 6)
    args.decoder_layers = getattr(args, "decoder_layers", 3)
    args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 6)
    mbart_base_architecture(args)

@register_model_architecture('bart', 'mbart_middle')
def mbart_middle(args):
    args.encoder_embed_dim = getattr(args, 'encoder_embed_dim', 768 // 2)
    args.encoder_ffn_embed_dim = getattr(args, 'encoder_ffn_embed_dim', 4*768 // 2)
    args.encoder_layers = getattr(args, 'encoder_layers', 6)
    args.encoder_attention_heads = getattr(args, 'encoder_attention_heads', 12 // 2)
    args.decoder_layers = getattr(args, 'decoder_layers', 6)
    args.decoder_attention_heads = getattr(args, 'decoder_attention_heads', 12 // 2)
