def build_model()

in codegen_sources/model/src/model/__init__.py [0:0]

• 69 lines of code
• 27 McCabe index (conditional complexity)

def build_model(params, dico, gpu=True):
    """
    Build model.
    """
    if params.encoder_only:
        # build
        model = TransformerModel(params, dico, is_encoder=True, with_output=True)

        # reload pretrained word embeddings
        if params.reload_emb != "":
            word2id, embeddings = load_embeddings(params.reload_emb, params)
            set_pretrain_emb(model, dico, word2id, embeddings, gpu)

        # reload a pretrained model
        if params.reload_model != "":
            logger.info("============ Model Reloading")
            logger.info("Reloading model from %s ..." % params.reload_model)
            reload_transformer(params, params.reload_model, dico, model, "model", gpu)

        logger.info("Model: {}".format(model))
        logger.info(
            "Number of parameters (model): %i"
            % sum([p.numel() for p in model.parameters() if p.requires_grad])
        )
        logger.info("")

        return [model.cuda() if gpu else model]

    else:
        # build
        # TODO: only output when necessary - len(params.clm_steps + params.mlm_steps) > 0
        encoder = TransformerModel(params, dico, is_encoder=True, with_output=True)

        if params.separate_decoders:
            decoders = [
                TransformerModel(params, dico, is_encoder=False, with_output=True)
                for _ in params.lang2id.values()
            ]
        else:
            decoders = [
                TransformerModel(params, dico, is_encoder=False, with_output=True)
            ]

        for layer in range(params.n_layers_decoder):
            if layer <= params.n_share_dec - 1:
                assert params.amp == -1, "sharing layers is not supported with AMP"
                logger.info("Sharing decoder attention parameters for layer %i" % layer)
                for i in range(1, len(decoders)):
                    decoders[i].attentions[layer] = decoders[0].attentions[layer]

        # reload pretrained word embeddings
        if params.reload_emb != "":
            word2id, embeddings = load_embeddings(params.reload_emb, params)
            set_pretrain_emb(encoder, dico, word2id, embeddings, gpu)
            for decoder in decoders:
                set_pretrain_emb(decoder, dico, word2id, embeddings, gpu)

        # reload a pretrained model
        if params.reload_model != "":
            logger.info("============ Model Reloading")
            enc_path, dec_path = params.reload_model.split(",")
            assert not (enc_path == "" and dec_path == "")

            # reload encoder
            if enc_path != "":
                logger.info("Reloading encoder from %s ..." % enc_path)
                reload_transformer(params, enc_path, dico, encoder, "encoder", gpu)

            # reload decoders
            if dec_path != "":
                for dec in decoders:
                    logger.info("Reloading decoders from %s ..." % dec_path)
                    if params.reload_encoder_for_decoder:
                        reload_transformer(params, dec_path, dico, dec, "encoder", gpu)
                    else:
                        reload_transformer(params, dec_path, dico, dec, "decoder", gpu)

        logger.debug("Encoder: {}".format(encoder))
        logger.debug("Decoder: {}".format(decoders))
        logger.info(
            "Number of parameters (encoder): %i"
            % sum([p.numel() for p in encoder.parameters() if p.requires_grad])
        )
        logger.info(
            "Number of parameters (decoders): %i"
            % sum([p.numel() for p in decoders[0].parameters() if p.requires_grad])
        )
        logger.info(f"Number of decoders: {len(decoders)}")
        logger.info("")

        return (
            [encoder.cuda() if gpu else encoder],
            [dec.cuda() if gpu else dec for dec in decoders],
        )