code: def extract_bert_weights(
      output_file,
      model_dir,
      head_num,
      pad_id=0,
  ):
      # load var names
      state_dict = torch.load(model_dir, "cpu")
  
      var_name_list = list(state_dict.keys())
  
      for name in var_name_list:
          if name.endswith("weight_quant._amax"):
              state_dict[name[:-12]] = torch.Tensor(
                  quantize(
                      state_dict[name[:-12]].numpy(),
                      127,
                      state_dict[name].numpy(),
                  )
              ).to(torch.uint8)
  
      # initialize output file
      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 var_name_list:
          name_split = name.split(".")
          if len(name_split) <= 3 or not name_split[3].isdigit():
              continue
          layer_id = int(name_split[3])
          enc_tensor_names.setdefault(layer_id, []).append(name)
  
      # fill encoder_stack
      for layer_id in sorted(enc_tensor_names.keys()):
          fill_quant_hdf5_layer(
              enc_tensor_names[layer_id],
              state_dict,
              hdf5_file,
              f"encoder_stack/{layer_id}/",
              enc_layer_mapping_dict,
          )
  
      # fill src_embedding - except for position embedding
      fill_quant_hdf5_layer(
          var_name_list,
          state_dict,
          hdf5_file,
          "src_embedding/",
          src_emb_mapping_dict,
      )
  
      # handling token_embeddings for BERT
      token_embedding = (
          state_dict["bert.embeddings.word_embeddings.weight"]
          + state_dict["bert.embeddings.token_type_embeddings.weight"][
              0
          ]
      )
      token_embedding = quantize(
          token_embedding.numpy(),
          127,
          state_dict[
              "bert.embeddings.emb_quant.clip.clip_value_max"
          ].numpy(),
      )
      print(
          f"processed token_embedding, shape: {token_embedding.shape}"
      )
      hdf5_file.create_dataset(
          "src_embedding/token_embedding",
          data=token_embedding,
          dtype="uint8",
      )
      hdf5_file.create_dataset(
          "src_embedding/emb_clip_max",
          data=state_dict[
              "bert.embeddings.emb_quant.clip.clip_value_max"
          ],
      )
  
      # 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/src_padding_id", data=pad_id, dtype="i4"
      )
      hdf5_file.create_dataset(
          "model_conf/is_post_ln", data=True, dtype="?"
      )
      hdf5_file.create_dataset(
          "model_conf/use_gelu", data=True, dtype="?"
      )
  
      # Move layernorm weights to match layernorm implementation in lightseq
      tmp_scale, tmp_bias = (
          hdf5_file["src_embedding/norm_scale"][()],
          hdf5_file["src_embedding/norm_bias"][()],
      )
      for layer_id in sorted(enc_tensor_names.keys()):
          new_tmp_scale = hdf5_file[
              f"encoder_stack/{layer_id}/multihead_norm_scale"
          ][()]
          new_tmp_bias = hdf5_file[
              f"encoder_stack/{layer_id}/multihead_norm_bias"
          ][()]
          hdf5_file[f"encoder_stack/{layer_id}/multihead_norm_scale"][
              ()
          ] = tmp_scale
          hdf5_file[f"encoder_stack/{layer_id}/multihead_norm_bias"][
              ()
          ] = tmp_bias
          tmp_scale, tmp_bias = new_tmp_scale, new_tmp_bias
  
          new_tmp_scale = hdf5_file[
              f"encoder_stack/{layer_id}/ffn_norm_scale"
          ][()]
          new_tmp_bias = hdf5_file[
              f"encoder_stack/{layer_id}/ffn_norm_bias"
          ][()]
          hdf5_file[f"encoder_stack/{layer_id}/ffn_norm_scale"][
              ()
          ] = tmp_scale
          hdf5_file[f"encoder_stack/{layer_id}/ffn_norm_bias"][
              ()
          ] = tmp_bias
          tmp_scale, tmp_bias = new_tmp_scale, new_tmp_bias
      hdf5_file["src_embedding/norm_scale"][()] = tmp_scale
      hdf5_file["src_embedding/norm_bias"][()] = tmp_bias
  
      hdf5_file.close()
      # read-in again to double check
      hdf5_file = h5py.File(output_file, "r")
  
      def _print_pair(key, value):
          if key == "sampling_method":
              value = "".join(map(chr, value[()]))
          else:
              value = value[()]
          print(f"{key}: {value}")
  
      list(
          map(
              lambda x: _print_pair(*x), hdf5_file["model_conf"].items()
          )
      )
  
file_path: lightseq/examples/inference/python/export/huggingface/ls_torch_hf_quant_bert_export.py
imports: import os
  import h5py
  from collections import OrderedDict
  import torch
  from lightseq.training.ops.pytorch.export import apply_rule
  from lightseq.training.ops.pytorch.export_quant import quantize
  from export.util import parse_args
label_desc: Use image node_id 2 for calling a global function with example usage: extract_bert_weights(output_file, model_dir, head_num, pad_id) without return types