def load_model(bison_args, device, data_handler, output_model_file=None):
"""
Load a model.
:param bison_args: instance of :py:class:BisonArguments
:param device: the device to move the model to
:param data_handler: the dataset handler, an instance of :py:class:BitextHandler or a subclass
:param output_model_file: the location of the model to load
:return: the loaded model
"""
model_state_dict = None
if output_model_file is not None:
model_state_dict = torch.load(output_model_file)
if bison_args.bert_model == 'bert-vanilla':
# randomly initialises BERT weights instead of using a pre-trained model
model = BertForMaskedLM(BertConfig.from_default_settings())
else:
model = BertForMaskedLM.from_pretrained(bison_args.bert_model,
state_dict=model_state_dict)
model.to(device)
return model
