def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, albert_config_file,
pytorch_dump_path):
# Initialise PyTorch model
config = AlbertConfig.from_json_file(albert_config_file)
print("Building PyTorch model from configuration: {}".format(str(config)))
model = AlbertForMaskedLM(config)
load_tf_weights_in_albert(model, config, tf_checkpoint_path)
print("Save PyTorch model to {}".format(pytorch_dump_path))
torch.save(model.state_dict(), pytorch_dump_path)
def albert_convert_tf_checkpoint_to_pytorch(tf_checkpoint_path,
albert_config_file,
pytorch_dump_path):
from transformers import AlbertConfig, AlbertForMaskedLM, load_tf_weights_in_albert
# Initialise PyTorch model
config = AlbertConfig.from_json_file(albert_config_file)
print("Building PyTorch model from configuration: {}".format(str(config)))
model = AlbertForMaskedLM(config)
# Load weights from tf checkpoint
load_tf_weights_in_albert(model, config, tf_checkpoint_path)
# Save pytorch-model
print("Save PyTorch model to {}".format(pytorch_dump_path))
torch.save(model.state_dict(), pytorch_dump_path)