code: def extract_vit_weights(
      output_file,
      model_dir,
      head_num,
      image_size,
      patch_size,
  ):
      # load var names
      encoder_state_dict = ViTModel.from_pretrained(
          model_dir
      ).state_dict()
  
      # Insert additional "final" to some weight to prevent ambiguous match
      def _insert_final(key):
          l = key.split(".")
          l.insert(3, "final")
          return ".".join(l)
  
      encoder_state_dict = OrderedDict(
          [
              (_insert_final(k), v)
              if len(k.split(".")) > 3 and k.split(".")[3] == "output"
              else (k, v)
              for k, v in encoder_state_dict.items()
          ]
      )
  
      enc_var_name_list = list(encoder_state_dict.keys())
  
      # initialize output file
      output_file += ".hdf5"
      print("Saving model to hdf5...")
      print("Writing to {0}".format(output_file))
      hdf5_file = h5py.File(output_file, "w")
  
      # fill each encoder layer's params
      enc_tensor_names = {}
      for name in enc_var_name_list:
          name_split = name.split(".")
          if len(name_split) <= 2 or not name_split[2].isdigit():
              continue
          layer_id = int(name_split[2])
          enc_tensor_names.setdefault(layer_id, []).append(name)
  
      # fill encoder_stack
      for layer_id in sorted(enc_tensor_names.keys()):
          fill_hdf5_layer(
              enc_tensor_names[layer_id],
              encoder_state_dict,
              hdf5_file,
              f"encoder_stack/{layer_id}/",
              enc_layer_mapping_dict,
          )
  
      # fill src_embedding - except for position embedding
      fill_hdf5_layer(
          enc_var_name_list,
          encoder_state_dict,
          hdf5_file,
          "src_embedding/",
          src_emb_mapping_dict,
      )
  
      # save number of layers metadata
      hdf5_file.create_dataset(
          "model_conf/n_encoder_stack",
          data=len(enc_tensor_names),
          dtype="i4",
      )
      # fill in model_conf
      hdf5_file.create_dataset(
          "model_conf/head_num", data=head_num, dtype="i4"
      )
      hdf5_file.create_dataset(
          "model_conf/use_gelu", data=True, dtype="?"
      )
      hdf5_file.create_dataset(
          "model_conf/is_post_ln", data=False, dtype="?"
      )
      hdf5_file.create_dataset(
          "model_conf/image_size", data=image_size, dtype="i4"
      )
      hdf5_file.create_dataset(
          "model_conf/patch_size", data=patch_size, dtype="i4"
      )
  
      hdf5_file.close()
      # read-in again to double check
      hdf5_file = h5py.File(output_file, "r")
  
      def _print_pair(key, value):
          value = value[()]
          print(f"{key}: {value}")
  
      list(
          map(
              lambda x: _print_pair(*x), hdf5_file["model_conf"].items()
          )
      )
  
imports: import os
  import h5py
  from collections import OrderedDict
  from transformers import ViTModel
  from lightseq.training.ops.pytorch.export import fill_hdf5_layer
label_desc: Use image node_id 1 for calling a global function with example usage: extract_vit_weights(output_file, model_dir, head_num, image_size, patch_size) and returns: str