code: def get_hf_models(model_name, config, dtype):
      pretrained_state_dict = state_dict_from_pretrained(model_name)
  
      def key_mapping_ln_gamma_beta(key):
          key = re.sub(r"LayerNorm.gamma$", "LayerNorm.weight", key)
          key = re.sub(r"LayerNorm.beta$", "LayerNorm.bias", key)
          return key
  
      pretrained_state_dict = OrderedDict(
          (key_mapping_ln_gamma_beta(k), v)
          for k, v in pretrained_state_dict.items()
      )
      model_hf = BertForPreTrainingHF(config)
      # Missing key(s) in state_dict: "bert.embeddings.position_ids", "cls.predictions.decoder.bias"
      # position_ids is a buffer, and predictions.decoder.bias is tied to predictions.bias.
      model_hf.load_state_dict(pretrained_state_dict, strict=False)
      model_hf.cuda().to(dtype=dtype)
      return model_hf
  
file_path: flash-attention/tests/models/test_bert.py
imports: import re
  from collections import OrderedDict
  import pytest
  import torch
  import torch.nn.functional
  from einops import rearrange
  from transformers import BertConfig
  from transformers.models.bert.modeling_bert import BertForPreTraining
  from transformers.models.bert.modeling_bert import BertModel
  from flash_attn.models.bert import BertForPreTraining, BertModel, inv_remap_state_dict, remap_state_dict
  from flash_attn.utils.pretrained import state_dict_from_pretrained
label_desc: Use image node_id 2 for calling a global function with example usage: get_hf_models(model_name, config, dtype) and returns: model_hf, key