Method __init__

SwissArmyTransformer/sat/model/official/mae_model.py:136–148

(self, args, transformer=None, layernorm_epsilon=1e-6)

class MAE(EncoderDecoderModel):
    def __init__(self, args, transformer=None, layernorm_epsilon=1e-6):
        encoder = MAEEncoder(args, transformer=transformer, layernorm_epsilon=layernorm_epsilon)
        dec_args = argparse.Namespace(**vars(args))
        # dec_args.enc_hidden_size = dec_args.hidden_size  # used for cross attn
        override_attrs = ['num_layers', 'hidden_size', 'num_attention_heads',
                            'max_sequence_length', 'inner_hidden_size', 'hidden_size_per_attention_head']
        for name in override_attrs:
            dec_attr = getattr(dec_args, 'dec_' + name, None)
            if dec_attr is not None:  # else use encoder-config
                setattr(dec_args, name, dec_attr)
        setattr(dec_args, 'enc_hidden_size', args.hidden_size)
        decoder = MAEDecoder(dec_args, transformer=transformer, layernorm_epsilon=layernorm_epsilon)
        super().__init__(args, encoder=encoder, decoder=decoder, tie_word_embeddings=False)
    
    def encode(self, input_ids, position_ids, attention_mask=None, **kw_args):
        return self.encoder(input_ids, position_ids, attention_mask, **kw_args)