def load(self) -> None:
if self.pretrained_bert:
log.info(f"From pretrained {self.pretrained_bert}.")
self.config = AutoConfig.from_pretrained(
self.pretrained_bert, num_labels=self.n_classes, output_attentions=True, output_hidden_states=True
)
self.model = BertModel.from_pretrained(self.pretrained_bert, config=self.config)
elif self.bert_config_file and Path(self.bert_config_file).is_file():
self.config = AutoConfig.from_json_file(str(expand_path(self.bert_config_file)))
self.model = BertModel.from_config(config=self.bert_config)
else:
raise ConfigError("No pre-trained BERT model is given.")
self.model.to(self.device)
