MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / load_hf_bert

Method load_hf_bert

tensorrt_llm/models/bert/model.py:131–147  ·  view source on GitHub ↗

Use as the abstractmethod, load corresponding HF model. Subclass must implement this method!

(cls, model_dir: str, load_model_on_cpu: bool,
                     dtype: torch.dtype)

Source from the content-addressed store, hash-verified

129
130 @classmethod
131 def load_hf_bert(cls, model_dir: str, load_model_on_cpu: bool,
132 dtype: torch.dtype):
133 """
134 Use as the abstractmethod, load corresponding HF model.
135 Subclass must implement this method!
136 """
137
138 assert cls.__name__ != "BertBase", f"Never call from BertBase class!"
139
140 if cls.__name__ == "BertModel":
141 return load_hf_bert_base(model_dir, load_model_on_cpu, dtype)
142 elif cls.__name__ == "BertForQuestionAnswering":
143 return load_hf_bert_qa(model_dir, load_model_on_cpu, dtype)
144 elif cls.__name__ == "BertForSequenceClassification":
145 return load_hf_bert_cls(model_dir, load_model_on_cpu, dtype)
146 else:
147 assert False, f"Unknown class {cls.__name__}!"
148
149 @classmethod
150 def from_hugging_face(

Callers 1

from_hugging_faceMethod · 0.80

Calls 3

load_hf_bert_baseFunction · 0.85
load_hf_bert_qaFunction · 0.85
load_hf_bert_clsFunction · 0.85

Tested by

no test coverage detected