from huggingface_hub import snapshot_download
import json
import os

REPO_ID = "mosaicml/mosaic-bert-base-seqlen-2048"
MODEL_DIRECTORY = "mosaic-bert-base-seqlen-2048"


def main():
    snapshot_download(repo_id=REPO_ID, local_dir=MODEL_DIRECTORY)

    # modify the model's config.json file to satisfy our requirements
    config_file_path = os.path.join(MODEL_DIRECTORY, 'config.json')
    contents = json.load(open(config_file_path))
    contents['architectures'] = ['BertModel']
    contents['auto_map']['AutoModel'] = 'bert_layers.BertModel'
    contents['torch_dtype'] = 'bfloat16'
    contents['transformers_version'] = '4.28.1'
    contents['_name_or_path'] = 'mosaic-bert-base-seqlen-2048'
    json.dump(contents, open(config_file_path, 'w'), ensure_ascii=True)

if __name__ == '__main__':
    main()