def _build_model_from_pretrain()

in src/hyperpod_nemo_adapter/collections/model/nlp/sagemaker_llama_model.py [0:0]

• 21 lines of code
• 3 McCabe index (conditional complexity)

    def _build_model_from_pretrain(self, model_cfg, torch_dtype=None, quantization_config=None):
        path = self._cfg.hf_model_name_or_path
        _logger.info("Loading pretrained weights from %s.", path)
        use_flash_attn = self._cfg.use_flash_attention
        access_token = self._cfg.get("hf_access_token", None)
        if TF_VERSION < pversion.parse("4.37.1") or not use_flash_attn:
            return Llama4ForConditionalGeneration.from_pretrained(
                path,
                config=model_cfg,
                torch_dtype=torch_dtype,
                quantization_config=quantization_config,
                token=access_token,
            )
        model_cfg.text_config._attn_implementation = "flash_attention_2"
        return Llama4ForConditionalGeneration.from_pretrained(
            path,
            config=model_cfg,
            torch_dtype=torch_dtype,
            quantization_config=quantization_config,
            token=access_token,
        )