code: def export_ls_embedding(
      file, state_dict, max_length, emb_dim, is_encoder, save_pb=True
  ):
      var_name_list = list(state_dict.keys())
      emb, target_tn = gather_token_embedding(
          var_name_list, state_dict, "embed_tokens", emb_dim=emb_dim
      )
      emb_size = emb.flatten().shape[0] - 1
      assert emb_size % emb_dim == 0
      emb = emb.flatten()[:-1].reshape(emb_size // emb_dim, emb_dim)
  
      if is_encoder:
          emb_list = emb.flatten().tolist()
          if save_pb:
              file.src_embedding.token_embedding[:] = emb_list
          else:
              file.create_dataset(
                  "src_embedding/token_embedding",
                  data=emb_list,
                  dtype="f4",
              )
      else:
          emb_list = emb.transpose().flatten().tolist()
          if save_pb:
              file.trg_embedding.token_embedding[:] = emb_list
          else:
              file.create_dataset(
                  "trg_embedding/token_embedding",
                  data=emb_list,
                  dtype="f4",
              )
      print(
          "%s -> %s_embedding.token_embedding, convert finished!"
          % (target_tn, "src" if is_encoder else "trg")
      )
  
      pos_emb = get_pos_embedding(max_length, emb_dim)
      pos_emb_list = pos_emb.flatten().tolist()
      if is_encoder:
          if save_pb:
              file.src_embedding.position_embedding[:] = pos_emb_list
          else:
              file.create_dataset(
                  "src_embedding/position_embedding",
                  data=pos_emb_list,
                  dtype="f4",
              )
      else:
          if save_pb:
              file.trg_embedding.position_embedding[:] = pos_emb_list
          else:
              file.create_dataset(
                  "trg_embedding/position_embedding",
                  data=pos_emb_list,
                  dtype="f4",
              )
      target_tn = [
          tn.replace("embeddings", "pos_embeddings") for tn in target_tn
      ]
      print(
          "%s -> %s_embedding.position_embedding, convert finished!"
          % (target_tn, "src" if is_encoder else "trg")
      )
  
file_path: lightseq/lightseq/training/ops/pytorch/export.py
imports: from collections import OrderedDict
  import numpy
  from lightseq.training.ops.pytorch.util import get_pos_embedding
  from lightseq.training import LSTransformerEncoderLayer, LSTransformerDecoderLayer
label_desc: Use image node_id 6 for calling a global function with example usage: export_ls_embedding(file, state_dict, max_length, emb_dim, is_encoder, save_pb) without return types