Method from_hugging_face

tensorrt_llm/models/bert/model.py:150–201

Create a BertModel object from give parameters

(
            cls,
            hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'],
            dtype: str = 'float16',
            mapping: Optional[Mapping] = None,
            quant_config: Optional[QuantConfig] = None,
            **kwargs)

    @classmethod
    def from_hugging_face(
            cls,
            hf_model_or_dir: Union[str, 'transformers.PreTrainedModel'],
            dtype: str = 'float16',
            mapping: Optional[Mapping] = None,
            quant_config: Optional[QuantConfig] = None,
            **kwargs):
        """
        Create a BertModel object from give parameters
        """
        import transformers

        assert hf_model_or_dir is not None
        use_preloading = isinstance(hf_model_or_dir,
                                    transformers.PreTrainedModel)
        if use_preloading:
            hf_model = hf_model_or_dir
            hf_config_or_dir = hf_model.config
        else:
            hf_model_dir = hf_model_or_dir
            hf_config_or_dir = hf_model_or_dir

        load_model_on_cpu = kwargs.pop('load_model_on_cpu', False)
        tllm_config = BERTConfig.from_hugging_face(
            hf_config_or_dir=hf_config_or_dir,
            dtype=dtype,
            mapping=mapping,
            quant_config=quant_config,
            **kwargs)
        #NOTE: override architecture info
        RobertaCls_mapping = {
            "BertModel": "RobertaModel",
            "BertForQuestionAnswering": "RobertaForQuestionAnswering",
            "BertForSequenceClassification": "RobertaForSequenceClassification",
        }
        if tllm_config.is_roberta:
            setattr(tllm_config, 'architecture',
                    RobertaCls_mapping[cls.__name__])
        else:
            setattr(tllm_config, 'architecture', cls.__name__)

        torch_dtype = torch.float16 if dtype == 'float16' else torch.float32
        if not use_preloading:
            hf_model = cls.load_hf_bert(model_dir=hf_model_dir,
                                        load_model_on_cpu=load_model_on_cpu,
                                        dtype=torch_dtype)
        weights = load_weights_from_hf_model(hf_model=hf_model,
                                             config=tllm_config)
        model = cls(tllm_config)
        model.load(weights)

        return model

    # Override the PretrainedModel's meothd, can unify in the future.
    def prepare_inputs(self, max_batch_size, max_input_len, **kwargs):