Upload llm-export/llm_export.py with huggingface_hub

llm-export/llm_export.py

@@ -0,0 +1,1467 @@
+import os
+import base64
+import glob
+import shutil
+import argparse
+import torch
+import numpy as np
+from onnxslim import slim
+import onnxruntime as ort
+import sentencepiece as spm
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, TaskType, get_peft_model, PeftModel
+try:
+    import _tools as MNNTools
+except:
+    MNNTools = None
+
+from llm_models.GOT.GOT_ocr_2_0 import GOTQwenForCausalLM
+from llm_models.GOT.modeling_qwen2 import Qwen2MLP, Qwen2RMSNorm, Qwen2Attention
+def onnx2mnn(onnx_path, mnn_dir, quant_bit = 4, asymmetric = True, external_data = False, bizCode : str= None):
+    model_name, model_extension = os.path.splitext(os.path.basename(onnx_path))
+    if model_extension != '.onnx':
+        return
+    mnn_name = model_name + '.mnn'
+    mnn_path = os.path.join(mnn_dir, mnn_name)
+    convert_args = [
+        '',
+        '-f',
+        'ONNX',
+        '--modelFile',
+        str(onnx_path),
+        '--MNNModel',
+        str(mnn_path),
+        '--weightQuantBits',
+        str(quant_bit),
+    ]
+    if asymmetric:
+        convert_args.append("--weightQuantAsymmetric")
+    if external_data:
+        convert_args.append("--saveExternalData")
+    if bizCode is not None:
+        convert_args.append("--bizCode")
+        convert_args.append(str(bizCode))
+    MNNTools.mnnconvert(convert_args)
+
+# some wrapper class for export
+class Embedding(torch.nn.Module):
+    def __init__(self, embed, using_bf16: bool = False):
+        super().__init__()
+        self.bf16 = using_bf16
+        self.embed_dim = embed.weight.shape[-1]
+        if using_bf16:
+            # using bf16 embedding weight
+            self.embed = embed.bfloat16()
+        else:
+            self.embed = embed
+
+    def forward(self, input_ids):
+        res = self.embed(input_ids)
+        if self.bf16:
+            res = res.float()
+        return res.view(-1, 1, self.embed_dim)
+
+class GOTEmbedding(torch.nn.Module):
+    def __init__(self, embed, using_bf16: bool = False):
+        super().__init__()
+        self.bf16 = using_bf16
+        self.embed_dim = embed.weight.shape[-1]
+        if using_bf16:
+            # using bf16 embedding weight
+            self.embed = embed.bfloat16()
+        else:
+            self.embed = embed
+
+    def forward(self, input_ids):
+        res = self.embed(input_ids)
+        if self.bf16:
+            res = res.float()
+        return res.view(1, -1, self.embed_dim)
+    
+class Lm(torch.nn.Module):
+    def __init__(self, lm):
+        super().__init__()
+        self.lm = lm
+
+    def forward(self, hidden_states):
+        m_logits = self.lm(hidden_states)
+        #token = torch.argmax(m_logits)
+        return m_logits
+
+class LLM(torch.nn.Module):
+    '''
+    Base class for all llm model. Inherits from [`torch.nn.Module`].
+    '''
+
+    def __init__(self, args):
+        super().__init__()
+        self.quant_bit = 4
+        self.asymmetric = True
+        self.onnx_path = args.onnx_path
+        self.mnn_path = args.mnn_path
+        if not os.path.exists(self.onnx_path):
+            os.makedirs(self.onnx_path)
+        if not os.path.exists(self.mnn_path):
+            os.makedirs(self.mnn_path)
+        self.export_mnn = args.export_mnn
+        self.export_verbose = args.export_verbose
+        self.export_test = args.export_test
+        # default is False, just set True when using below command:
+        # `python llm_export ../path --export --embed_bin` to export single model without embedding
+        self.without_embed = False
+        self.embed_bin = args.embed_bin
+        if self.embed_bin:
+            self.embed_bf16 = True
+        else:
+            self.embed_bf16 = args.embed_bf16
+        self.skip_slim = args.skip_slim
+        tokenizer_model = os.path.join(args.path, 'tokenizer.model')
+        if os.path.exists(tokenizer_model):
+            self.sp_model = spm.SentencePieceProcessor(tokenizer_model)
+        else:
+            self.sp_model = None
+        merge_file = os.path.join(args.path, 'merges.txt')
+        if os.path.exists(merge_file):
+            self.merge_txt = merge_file
+        else:
+            self.merge_txt = None
+        self.stop_ids = []
+        self.max_length = 1024
+        self.hidden_size = 4096
+        self.visual = None # defualt is not visual
+        self.lora_path = args.lora_path
+        self.load_hf(args.path)
+        self.load_model()
+
+    def load_hf(self, model_path: str):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).float().eval()
+        self.config = self.model.config
+        if self.lora_path is not None:
+            adapter = PeftModel.from_pretrained(self.model, model_id=self.lora_path)
+            self.model = adapter.merge_and_unload(progressbar=True)
+
+    def load_model(self):
+        raise NotImplementedError
+
+    def get_attention_mask(self) -> torch.Tensor:
+        raise NotImplementedError
+
+    def get_position_ids(self) -> torch.Tensor:
+        raise NotImplementedError
+
+    def export_vocab(self):
+        raise NotImplementedError
+
+    def visual_embed(self, input_ids):
+        raise NotImplementedError
+
+    def __embedding(self, input_ids):
+        if self.visual is not None and self.token_len == 0:
+            input_embeds = self.visual_embed(input_ids)
+        else:
+            input_embeds = self.embed(input_ids)
+        return input_embeds
+
+    def __decode(self, hidden_states, attention_mask, position_ids, past_key_values):
+        presents = []
+        for i in range(self.block_nums):
+            hidden_states, kv = self.blocks[i](hidden_states, attention_mask, position_ids, past_key_values[i])
+            presents.append(kv)
+        token_id = self.lm(hidden_states).view(1)
+        presents = torch.stack(presents)
+        self.seq_len += 1
+        self.token_len += 1
+        return token_id, presents
+
+    def forward(self, input_ids, attention_mask, position_ids, past_key_values):
+        if self.without_embed:
+            return self.__decode(input_ids, attention_mask, position_ids, past_key_values)
+        return self.__decode(self.__embedding(input_ids), attention_mask, position_ids, past_key_values)
+
+    # some test functions
+    def build_prompt(self, query):
+        if hasattr(self.tokenizer, 'build_prompt'):
+            prompt = self.tokenizer.build_prompt(query)
+        else:
+            prompt = query
+        return prompt
+
+    def str_to_ids(self, prompt):
+        input_ids = self.tokenizer(prompt, return_tensors="pt")['input_ids']
+        return input_ids
+
+    def id_to_str(self, token_id):
+        word = self.tokenizer._convert_id_to_token(int(token_id))
+        word = self.tokenizer.convert_tokens_to_string([word])
+        return word
+
+    def response(self, query):
+        prompt = self.build_prompt(query)
+        input_ids = self.str_to_ids(prompt)
+        self.seq_len = input_ids.numel()
+        self.context_len = self.seq_len - 2
+        self.token_len = 0
+        past_key_values = [None for i in range(self.block_nums)]
+        token_id = input_ids
+        while self.token_len < self.max_length:
+            attention_mask = self.get_attention_mask()
+            position_ids = self.get_position_ids()
+            token_id, past_key_values = self.forward(token_id, attention_mask, position_ids, past_key_values)
+            if token_id == self.stop_id or token_id in self.stop_ids:
+                print("", end='\n')
+                break
+            word = self.id_to_str(token_id)
+            print(word, end="", flush=True)
+
+    # some export functions
+    def assert_equal(self, torch_outs, onnx_outs):
+        if type(torch_outs) not in (list, tuple):
+            torch_outs = (torch_outs, )
+            onnx_outs = (onnx_outs, )
+        same = True
+        for orig, onnx in zip(torch_outs, onnx_outs):
+            orig = orig.detach().numpy()
+            if not np.allclose(orig, onnx, rtol=1e-3, atol=1e-3):
+                print('Error: onnx outputs dont match original. [shape = {}] onnx: {}, original: {}'.format(onnx.shape, onnx, orig))
+                same = False
+                break
+        if same:
+            print('onnx test SUCCESS')
+
+    def export_lm(self):
+        model = self.lm
+        hidden_states = torch.randn(1, self.hidden_size)
+        onnx_model = f'./{self.onnx_path}/lm.onnx'
+        torch.onnx.export(model, (hidden_states),
+                        onnx_model,
+                        verbose=self.export_verbose,
+                        input_names=['hidden_states'],
+                        output_names=['token_id'],
+                        do_constant_folding=True,
+                        opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+        # test lm
+        if self.export_test:
+            original_outs = model(hidden_states)
+            ort_session = ort.InferenceSession(onnx_model, providers=['CPUExecutionProvider'])
+            inputs = {
+                'hidden_states' : hidden_states.numpy(),
+            }
+            onnx_outs = ort_session.run(None, inputs)
+            self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path, self.quant_bit, self.asymmetric)
+
+    def export_visual(self):
+        if self.visual is None:
+            return
+        input_images = torch.randn((1, 3, self.image_size, self.image_size))
+        model = self.visual
+        onnx_model = f'./{self.onnx_path}/visual.onnx'
+        torch.onnx.export(model, (input_images),
+                        onnx_model,
+                        verbose=self.export_verbose,
+                        input_names=['input_images'],
+                        output_names=['image_embeds'],
+                        dynamic_axes={"input_images": {
+                            0: "size"
+                        }},
+                        do_constant_folding=True,
+                        opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+        # test
+        if self.export_test:
+            original_outs = model(input_images)
+            ort_session = ort.InferenceSession(onnx_model, providers=['CPUExecutionProvider'])
+            inputs = {
+                'input_images' : input_images.numpy(),
+            }
+            onnx_outs = ort_session.run(None, inputs)[0]
+            self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
+
+    def export_embed(self):
+        model = self.embed
+        if self.embed_bin:
+            import ctypes
+            tensor_data = model.embed.weight.data
+            data_ptr = tensor_data.untyped_storage().data_ptr()
+            buffer = (ctypes.c_byte * (tensor_data.numel() * 2)).from_address(data_ptr)
+            with open(f'./{self.mnn_path}/embeddings_bf16.bin', 'wb') as f:
+                f.write(buffer)
+            return
+        input_ids = torch.arange(3, dtype=torch.long)
+        onnx_model = f'./{self.onnx_path}/embedding.onnx'
+        torch.onnx.export(model, (input_ids),
+                        onnx_model,
+                        verbose=self.export_verbose,
+                        input_names=['input_ids'],
+                        output_names=['inputs_embeds'],
+                        dynamic_axes={"input_ids": {
+                            0: "length"
+                        }},
+                        do_constant_folding=True,
+                        opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+        # test
+        if self.export_test:
+            original_outs = model(input_ids)
+            ort_session = ort.InferenceSession(onnx_model, providers=['CPUExecutionProvider'])
+            inputs = {
+                'input_ids' : input_ids.numpy(),
+            }
+            onnx_outs = ort_session.run(None, inputs)
+            self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
+
+    def export_block(self, block_id: int):
+        self.seq_len = 3
+        self.token_len = 0
+        inputs_embeds = torch.randn((1, self.seq_len, self.hidden_size))
+        attention_mask =  self.get_attention_mask()
+        position_ids = self.get_position_ids()
+        past_key_cache = torch.randn((1,  self.num_key_value_heads, 0, self.hidden_size// self.num_key_value_heads))  # torch.Size([1, 16, 286, 64])
+        past_value_cache = torch.randn((1, self.num_key_value_heads, 0, self.hidden_size// self.num_key_value_heads))
+        model = self.blocks[block_id]
+        onnx_model = f'./{self.onnx_path}/block_{block_id}.onnx'
+        # 每一个 循环都有pastkv cache
+        torch.onnx.export(
+            model, (inputs_embeds, attention_mask, position_ids,past_key_cache,past_value_cache),
+            onnx_model,
+            verbose=self.export_verbose,
+            input_names=[
+                'inputs_embeds', 'attention_mask', 'position_ids', 'past_key_cache', 'past_value_cache'
+            ],
+            output_names=['hidden_states', 'past_key_states', 'past_value_states'],
+            dynamic_axes= {
+                            "inputs_embeds" : { 1: "seq_len" },
+                            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+                            "position_ids" : { 1: "seq_len" },
+                            "past_key_cache" : { 2: "seq_len" },
+                            "past_value_cache" : { 2: "seq_len" },
+                            "hidden_states":{1: "seq_len" },
+                            "past_key_states":{2: "seq_len" },
+                            "past_value_states":{2: "seq_len" },
+                        },
+            opset_version=17)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+        if self.export_test:
+            original_outs = model(inputs_embeds, attention_mask, position_ids)
+            ort_session = ort.InferenceSession(onnx_model, providers=['CPUExecutionProvider'])
+            inputs = {
+                'inputs_embeds' : inputs_embeds.detach().numpy(),
+                'attention_mask' : attention_mask.numpy(),
+                'position_ids' : position_ids.numpy(),
+            }
+            onnx_outs = ort_session.run(None, inputs)
+            self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path, self.quant_bit, self.asymmetric)
+
+    def export_blocks(self):
+        for i in range(self.block_nums):
+            self.export_block(i)
+
+    def export(self):
+        model = self
+        self.seq_len = 3
+        self.token_len = 0
+        input_ids = torch.arange(3, dtype=torch.long)
+        attention_mask =  self.get_attention_mask()
+        position_ids = self.get_position_ids()
+        past_key_values = torch.zeros(self.past_kv_shape)
+        onnx_model = f'./{self.onnx_path}/llm.onnx'
+        if self.embed_bin:
+            self.without_embed = True
+            input_ids = self.__embedding(input_ids)
+        print('export start ...')
+        torch.onnx.export(
+            model, (input_ids, attention_mask, position_ids, past_key_values),
+            onnx_model,
+            verbose=self.export_verbose,
+            input_names=[
+                'input_ids', 'attention_mask', 'position_ids', 'past_key_values'
+            ],
+            output_names=['token_id', 'presents'],
+            dynamic_axes=self.model_dynamic_axes,
+            do_constant_folding=True,
+            opset_version=15)
+        print('export done!')
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+        if self.export_test:
+            # test
+            original_outs = model(input_ids, attention_mask, position_ids, past_key_values)
+            ort_session = ort.InferenceSession(onnx_model, providers=['CPUExecutionProvider'])
+            inputs = {
+                'input_ids' : input_ids.detach().numpy(),
+                'attention_mask' : attention_mask.numpy(),
+                'position_ids' : position_ids.numpy(),
+                'past_key_values' : past_key_values.numpy()
+            }
+            onnx_outs = ort_session.run(None, inputs)
+            self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            # single model is > 2G, using external_data
+            onnx2mnn(onnx_model, self.mnn_path, self.quant_bit, self.asymmetric, True)
+        if self.without_embed:
+            self.without_embed = False
+
+    def export_tokenizer(self):
+        file_path = os.path.join(self.onnx_path, "tokenizer.txt")
+        if self.sp_model is not None:
+            # senetencepiece
+            print('# senetencepiece tokenier')
+            NORMAL = 1; UNKNOWN = 2; CONTROL = 3
+            USER_DEFINED = 4; UNUSED = 5; BYTE = 6
+            fp = open(file_path, "w", encoding="utf8")
+            for i in range(self.sp_model.GetPieceSize()):
+                token = self.sp_model.IdToPiece(i)
+                score = self.sp_model.GetScore(i)
+                type = NORMAL
+                if self.sp_model.IsUnknown(i):
+                    type = UNKNOWN
+                elif self.sp_model.IsControl(i):
+                    type = CONTROL
+                elif self.sp_model.IsUnused(i):
+                    type = UNUSED
+                elif self.sp_model.IsByte(i):
+                    type = BYTE
+                if self.model_name == 'Chatglm_6b':
+                    if '<n>' in token: token = '\n'
+                    if '<|tab|>' in token: token = '\t'
+                    if '<|blank_' in token: token = ' ' * int(token[8:token.find('|>')])
+                if '▁' in token: token = token.replace('▁', ' ')
+                token_encode = base64.b64encode(token.encode("utf-8")).decode("utf8")
+                fp.write(f'{token_encode} {score} {type}\n')
+            fp.close()
+        elif hasattr(self.tokenizer, 'mergeable_ranks'):
+            print('# tiktoken tokenier')
+            # tikton
+            with open(file_path, "w", encoding="utf8") as fp:
+                for k, v in self.tokenizer.mergeable_ranks.items():
+                    line = base64.b64encode(k).decode("utf8") + "\n"
+                    fp.write(line)
+                if hasattr(self.tokenizer, 'special_tokens'):
+                    for k, v in self.tokenizer.special_tokens.items():
+                        line = base64.b64encode(k.encode("utf-8")).decode("utf8") + "\n"
+                        fp.write(line)
+        elif self.merge_txt is not None:
+            # huggingface tokenizer
+            merge_list = []
+            vocab = self.tokenizer.get_vocab()
+            vocab_list = ['<unk>' for i in range(len(vocab))]
+            # load vocab
+            for k, v in vocab.items():
+                vocab_list[int(v)] = k
+            # load merge
+            with open(self.merge_txt, 'rt') as merge:
+                for line in merge.readlines():
+                    merge_list.append(line)
+            # write to tokenizer.txt
+            with open(file_path, "w", encoding="utf8") as fp:
+                fp.write(f'{len(vocab_list)} {len(merge_list)}\n')
+                for v in vocab_list:
+                    fp.write(v + '\n')
+                for m in merge_list:
+                    fp.write(m)
+        else:
+            # huggingface tokenizer
+            def unicode_to_byte(u: int):
+                if u >= 256 and u <= 288:
+                    return u - 256
+                if u >= 289 and u <= 322:
+                    return u - 162
+                if u == 323:
+                    return 173
+                if u == 65372: # |
+                    return 124
+                if u == 9601:  # _
+                    return 95
+                return u
+            with open(file_path, "w", encoding="utf8") as fp:
+                vocab = self.tokenizer.get_vocab()
+                vocab_list = ['<unk>' for i in range(len(vocab))]
+                for k, v in vocab.items():
+                    try:
+                        vocab_list[int(v)] = bytes([unicode_to_byte(ord(c)) for c in k]).decode('utf-8', errors='ignore')
+                    except:
+                        vocab_list[int(v)] = k
+                for v in vocab_list:
+                    line = base64.b64encode(v.encode('utf-8')).decode("utf8") + "\n"
+                    fp.write(line)
+
+
+# chatglm
+class GLMBlock(torch.nn.Module):
+    def __init__(self, block, block_id, final_layernorm = None):
+        super().__init__()
+        self.block = block
+        self.block_id = block_id
+        self.final_layernorm = final_layernorm
+
+    def forward(self, hidden_states, attention_mask, position_ids, past_kv):
+        hidden_states, presents = self.block(hidden_states,
+                                             position_ids,
+                                             attention_mask,
+                                             self.block_id,
+                                             past_kv,
+                                             use_cache=True)
+        if self.final_layernorm is not None:
+            hidden_states = self.final_layernorm(hidden_states)
+            hidden_states = hidden_states.view(-1, self.hidden_size)[-1].view(1, 1, self.hidden_size)
+        if isinstance(presents, tuple):
+            presents = torch.stack(presents)
+        return hidden_states, presents
+
+class Chatglm_6b(LLM):
+    def __init__(self, args):
+        super().__init__(args)
+        self.model_name = 'Chatglm_6b'
+
+    def load_model(self):
+        transformer = self.model.transformer
+        self.lm_ = self.model.lm_head
+        self.embed_ = transformer.word_embeddings
+        self.blocks_ = transformer.layers
+        self.final_layernorm_ = transformer.final_layernorm
+        # some wrapper
+        self.stop_id = self.tokenizer._convert_token_to_id(self.tokenizer.eos_token)
+        self.block_nums = len(self.blocks_)
+        self.lm = Lm(self.lm_)
+        # chatglm embedding and lm using same param, copy embedding when using bf16
+        if self.embed_bf16:
+            import copy
+            embed_copy = copy.deepcopy(self.embed_)
+            self.embed = Embedding(embed_copy, self.embed_bf16)
+        else:
+            self.embed = Embedding(self.embed_, self.embed_bf16)
+        self.blocks = [GLMBlock(self.blocks_[i], i, self.final_layernorm_ if i == len(self.blocks_) - 1 else None) for i in range(self.block_nums)]
+        # some config for export
+        self.past_kv_shape = [28, 2, 0, 1, 32, 128]
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 2: "seq_len" },
+            "past_key_values" : { 1: "history_len" }
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 2: "seq_len" },
+            "past_key_values" : { 2: "history_len" }
+        }
+
+    def get_attention_mask(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.zeros([1]).bool().reshape([1, 1, 1, 1])
+        attention_mask = torch.zeros([self.seq_len, self.seq_len], dtype=torch.bool)
+        for i in range(self.seq_len):
+            attention_mask[i][-1] = True
+        attention_mask = attention_mask.reshape([1, 1, self.seq_len, self.seq_len])
+        return attention_mask
+
+    def get_position_ids(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.tensor([1, self.seq_len - self.context_len]).reshape([1, 2, 1])
+        position_ids_0 = torch.arange(self.seq_len, dtype=torch.long)
+        position_ids_1 = torch.zeros(self.seq_len, dtype=torch.long)
+        position_ids_1[-1] = 1
+        position_ids = torch.stack([position_ids_0, position_ids_1]).view(1, 2, -1)
+        return position_ids
+
+# chatglm2
+class GLM2Block(torch.nn.Module):
+    def __init__(self, block, block_id, final_layernorm = None):
+        super().__init__()
+        self.block = block
+        self.block_id = block_id
+        self.final_layernorm = final_layernorm
+        self.hidden_size = 4096
+
+    def forward(self, hidden_states, attention_mask, position_ids, past_kv):
+        theta = 1.0 / (10000 ** (torch.arange(0, 64, 2, dtype=torch.float32) / 64))
+        position_ids = position_ids.float().reshape(-1, 1)
+        idx_theta = position_ids * theta
+        rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1).unsqueeze(0).contiguous()
+        hidden_states, presents = self.block(hidden_states,
+                                            attention_mask,
+                                            kv_cache=past_kv,
+                                            rotary_pos_emb=rotary_pos_emb)
+        if self.final_layernorm is not None:
+            hidden_states = self.final_layernorm(hidden_states)
+            hidden_states = hidden_states.view(-1, self.hidden_size)[-1].view(1, 1, self.hidden_size)
+        if isinstance(presents, tuple):
+            presents = torch.stack(presents)
+        return hidden_states, presents
+
+class Chatglm2_6b(LLM):
+    def __init__(self, args):
+        super().__init__(args)
+        self.model_name = 'Chatglm2_6b'
+        if 'codegeex2-6b' in args.path:
+            self.model_name = 'Codegeex2_6b'
+
+    def load_model(self):
+        transformer = self.model.transformer
+        self.lm_ = transformer.output_layer
+        self.embed_ = transformer.embedding.word_embeddings
+        self.blocks_ = transformer.encoder.layers
+        self.final_layernorm_ = transformer.encoder.final_layernorm
+        # some wrapper
+        self.stop_id = self.tokenizer.eos_token_id
+        if self.stop_id is None:
+            # codegeex2-6b
+            self.stop_id = self.tokenizer.tokenizer.eos_id
+        self.block_nums = len(self.blocks_)
+        self.embed = Embedding(self.embed_, self.embed_bf16)
+        self.lm = Lm(self.lm_)
+        self.blocks = [GLM2Block(self.blocks_[i], i, self.final_layernorm_ if i == len(self.blocks_) - 1 else None) for i in range(self.block_nums)]
+        # some config for export
+        self.past_kv_shape = [28, 2, 0, 1, 2, 128]
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 1: "history_len" }
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 2: "history_len" }
+        }
+
+    def get_attention_mask(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.zeros([1, 1, 1, 1]).bool()
+        attention_mask = ~torch.tril(torch.ones([1, 1, self.seq_len, self.seq_len]).bool())
+        return attention_mask
+
+    def get_position_ids(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.tensor([self.token_len], dtype=torch.long)
+        return torch.arange(self.seq_len, dtype=torch.long)
+
+# chatglm3
+class Chatglm3_6b(Chatglm2_6b):
+    def __init__(self, args):
+        super().__init__(args)
+        self.model_name = 'Chatglm3_6b'
+
+    def build_prompt(self, query):
+        return f'<|user|>\n{query}\n<|assistant|>\n'
+
+# qwen
+class QWENBlock(torch.nn.Module):
+    def __init__(self, name, block, block_id, hidden_size, final_layernorm = None):
+        super().__init__()
+        self.name = name
+        self.block = block
+        self.block_id = block_id
+        self.final_layernorm = final_layernorm
+        self.hidden_size = hidden_size
+
+    def forward(self, hidden_states, attention_mask, position_ids, past_kv):
+        theta = 1.0 / (10000.0 ** (torch.arange(0, 128, 2, dtype=torch.float32) / 128))
+        position_ids = position_ids.float().reshape(-1, 1)
+        idx_theta = position_ids * theta
+        rotary_pos_emb = torch.cat((idx_theta, idx_theta), dim=-1)
+        rotary_pos_emb = rotary_pos_emb.unsqueeze(1).unsqueeze(0)
+        if self.name != 'Qwen-7B':
+            rotary_pos_emb = torch.stack([torch.cos(rotary_pos_emb), torch.sin(rotary_pos_emb)])
+        hidden_states = hidden_states.view(1, -1, self.hidden_size)
+        hidden_states, presents = self.block(hidden_states=hidden_states,
+                                             layer_past=past_kv,
+                                             attention_mask=attention_mask,
+                                             rotary_pos_emb=rotary_pos_emb,
+                                             use_cache=True)
+        if self.final_layernorm is not None:
+            hidden_states = self.final_layernorm(hidden_states)
+            hidden_states = hidden_states.view(-1, self.hidden_size)[-1].view(1, 1, self.hidden_size)
+        if isinstance(presents, tuple):
+            presents = torch.stack(presents)
+        return hidden_states, presents
+
+class QWEN18Block(torch.nn.Module):
+    def __init__(self, block, block_id, hidden_size, final_layernorm = None):
+        super().__init__()
+        self.block = block
+        self.block_id = block_id
+        self.final_layernorm = final_layernorm
+        self.hidden_size = hidden_size
+
+    def forward(self, hidden_states, attention_mask, position_ids, past_kv):
+        theta = 1.0 / (10000.0 ** (torch.arange(0, 128, 2, dtype=torch.float32) / 128))
+        position_ids = position_ids.float().reshape(-1, 1)
+        idx_theta = position_ids * theta
+        rotary_pos_emb = torch.cat((idx_theta, idx_theta), dim=-1).unsqueeze(1).unsqueeze(0)
+        rotary_pos_emb = torch.stack([torch.cos(rotary_pos_emb), torch.sin(rotary_pos_emb)])
+        hidden_states = hidden_states.view(1, -1, self.hidden_size)
+        hidden_states, presents = self.block(hidden_states,
+                                             rotary_pos_emb,
+                                             past_kv,
+                                             attention_mask,
+                                             use_cache=True)
+        if self.final_layernorm is not None:
+            hidden_states = self.final_layernorm(hidden_states)
+            hidden_states = hidden_states.view(-1, self.hidden_size)[-1].view(1, 1, self.hidden_size)
+        if isinstance(presents, tuple):
+            presents = torch.stack(presents)
+        return hidden_states, presents
+
+class Qwen_Chat(LLM):
+    def __init__(self, args):
+        super().__init__(args)
+
+    def load_model(self):
+        # Qwen models
+        self.model_name = 'Qwen-7B'
+        if '1_8' in model_path:
+            self.model_name = 'Qwen-1_8b'
+        if 'VL' in model_path:
+            self.model_name = 'Qwen-VL'
+        transformer = self.model.transformer
+        self.lm_ = self.model.lm_head
+        self.embed_ = transformer.wte
+        self.blocks_ = transformer.h
+        self.final_layernorm_ = transformer.ln_f
+        if hasattr(transformer, 'visual'):
+            self.visual = transformer.visual
+            self.image_start_id = transformer.config.visual['image_start_id']
+            self.image_size = transformer.config.visual['image_size']
+        # some wrapper
+        self.stop_id = self.tokenizer.im_end_id
+        self.block_nums = len(self.blocks_)
+        self.hidden_size = transformer.embed_dim
+        self.embed = Embedding(self.embed_, self.embed_bf16)
+        self.lm = Lm(self.lm_)
+        self.blocks = [QWENBlock(self.model_name, self.blocks_[i], i, self.hidden_size, self.final_layernorm_ if i == len(self.blocks_) - 1 else None) for i in range(self.block_nums)]
+        if self.block_nums == 32:
+            # qwen-7b, qwen-vl
+            self.past_kv_shape = [32, 2, 1, 0, 32, 128]
+        elif self.block_nums == 24:
+            # qwen-1.8b
+            self.past_kv_shape = [24, 2, 1, 0, 16, 128]
+        # some config for export
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 2: "history_len" }
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 3: "history_len" }
+        }
+
+    def build_prompt(self, query):
+        return f'\n<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n'
+
+    def get_attention_mask(self) -> torch.Tensor:
+        if self.model_name == 'Qwen-VL':
+            if self.token_len:
+                return torch.zeros([1, 1, 1, self.seq_len], dtype=torch.float32)
+            return (1 - torch.tril(torch.ones([1, 1, self.seq_len, self.seq_len]))) * torch.finfo(torch.float32).min
+        if self.token_len:
+            return torch.ones([1, 1, 1, 1]).bool()
+        return torch.tril(torch.ones([1, 1, self.seq_len, self.seq_len]).bool())
+
+    def get_position_ids(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.tensor([self.seq_len - 1], dtype=torch.long)
+        return torch.arange(self.seq_len, dtype=torch.long)
+
+    def visual_embed(self, input_ids):
+        if not torch.any(input_ids == self.image_start_id):
+            return self.embed(input_ids)
+        bos_pos = torch.where(input_ids == self.image_start_id)
+        eos_pos = torch.where(input_ids == self.image_start_id + 1)
+        img_pos = torch.stack((bos_pos[0], bos_pos[1], eos_pos[1]), dim=1)
+        images = []
+        for i, a, b in img_pos:
+            image = input_ids[i][a + 1 : b - 1].tolist()
+            image = image[ : image.index(self.image_start_id + 2)]
+            images.append(bytes(image).decode('utf-8'))
+        images = self.visual.encode(images)
+        hidden_states = self.embed(input_ids).view(1, -1, self.hidden_size)
+        for idx, (i, a, b) in enumerate(img_pos):
+            hidden_states[i][a + 1 : b] = images[idx]
+        return hidden_states.view(-1, 1, self.hidden_size)
+
+
+class Qwen2DecoderLayer(torch.nn.Module):
+    def __init__(self, config, block, layer_idx: int):
+        super().__init__()
+        self.block = block
+        # self.hidden_size = config.hidden_size
+        self.self_attn = Qwen2Attention(config, layer_idx)
+        # 加载权重
+        self.self_attn.load_state_dict(block.self_attn.state_dict())
+        self.mlp = self.block.mlp
+        self.input_layernorm = self.block.input_layernorm
+        self.post_attention_layernorm = self.block.post_attention_layernorm
+        # self.mlp = Qwen2MLP(config)
+        # self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask,
+        position_ids,
+        past_key_cache=None,
+        past_value_cache=None
+    ):
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states,past_key_states,past_value_states = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_cache=past_key_cache,
+            past_value_cache=past_value_cache
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        return hidden_states,past_key_states,past_value_states
+        #hidden_states = self.block(hidden_states, attention_mask, position_ids)
+        #return hidden_states
+
+
+class QWEN2Block(torch.nn.Module):
+    def __init__(self, name, block, block_id, config, final_layernorm = None):
+        super().__init__()
+        self.name = name
+        self.block = block
+        self.block_id = block_id
+        self.final_layernorm = final_layernorm
+        self.hidden_size = config.hidden_size
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.rope_theta = config.rope_theta
+
+    def forward(self, hidden_states, attention_mask, position_ids):
+        theta = 1.0 / (self.rope_theta ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim))
+        position_ids = position_ids.float().reshape(-1, 1)
+        idx_theta = position_ids * theta
+        rotary_pos_emb = torch.cat((idx_theta, idx_theta), dim=-1)
+        rotary_pos_emb = rotary_pos_emb.unsqueeze(1).unsqueeze(0)
+        rotary_pos_emb = torch.stack([torch.cos(rotary_pos_emb), torch.sin(rotary_pos_emb)])
+        hidden_states = hidden_states.view(1, -1, self.hidden_size)
+        hidden_states = self.block(hidden_states=hidden_states,
+                                             attention_mask=attention_mask,
+                                             #past_key_value=past_kv,
+                                             rotary_pos_emb=rotary_pos_emb,
+                                             #use_cache=True
+                                             )
+        if self.final_layernorm is not None:
+            hidden_states = self.final_layernorm(hidden_states[0])
+            hidden_states = hidden_states.view(-1, self.hidden_size)[-1].view(1, 1, self.hidden_size)
+        # print('###', presents.shape)
+        return hidden_states
+
+class Qwen2_Chat(LLM):
+    def __init__(self, args):
+        super().__init__(args)
+
+    def load_model(self):
+        # Qwen2 models
+        self.model_name = 'Qwen2'
+        transformer = self.model.model
+        self.lm_ = self.model.lm_head
+        self.embed_ = transformer.embed_tokens
+        self.blocks_ = transformer.layers
+        self.final_layernorm_ = transformer.norm
+        # some wrapper
+        self.stop_id = self.tokenizer.eos_token_id
+        if hasattr(model, 'generation_config'):
+            self.stop_ids.append(self.stop_id)
+            for id in self.model.generation_config.eos_token_id:
+                self.stop_ids.append(id)
+        self.block_nums = self.config.num_hidden_layers
+        self.hidden_size = self.config.hidden_size
+        self.num_heads = self.config.num_attention_heads
+        self.rope_theta = self.config.rope_theta
+        self.head_dim = self.hidden_size // self.num_heads
+        if self.embed_.weight is self.lm_.weight:
+            import copy
+            embed_copy = copy.deepcopy(self.embed_)
+            self.embed = Embedding(embed_copy, self.embed_bf16)
+        else:
+            self.embed = Embedding(self.embed_, self.embed_bf16)
+        self.lm = Lm(self.lm_)
+        self.past_kv_shape = [self.block_nums, 2, 1, 0, self.num_heads, self.head_dim]
+        self.blocks = [QWEN2Block(self.model_name, self.blocks_[i], i, self.config, self.final_layernorm_ if i == len(self.blocks_) - 1 else None) for i in range(self.block_nums)]
+        # some config for export
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 1: "history_len" }
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 2: "history_len" }
+        }
+
+    def build_prompt(self, query):
+        return f'<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n'
+
+    def get_attention_mask(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.zeros([1, 1, 1, self.seq_len], dtype=torch.float32)
+        return (1 - torch.tril(torch.ones([1, 1, self.seq_len, self.seq_len]))) * torch.finfo(torch.float32).min
+
+
+    def get_position_ids(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.tensor([[self.seq_len - 1]], dtype=torch.long)
+        return torch.arange(self.seq_len, dtype=torch.long).unsqueeze(0)
+
+    def visual_embed(self, input_ids):
+        if not torch.any(input_ids == self.image_start_id):
+            return self.embed(input_ids)
+        bos_pos = torch.where(input_ids == self.image_start_id)
+        eos_pos = torch.where(input_ids == self.image_start_id + 1)
+        img_pos = torch.stack((bos_pos[0], bos_pos[1], eos_pos[1]), dim=1)
+        images = []
+        for i, a, b in img_pos:
+            image = input_ids[i][a + 1 : b - 1].tolist()
+            image = image[ : image.index(self.image_start_id + 2)]
+            images.append(bytes(image).decode('utf-8'))
+        images = self.visual.encode(images)
+        hidden_states = self.embed(input_ids).view(1, -1, self.hidden_size)
+        for idx, (i, a, b) in enumerate(img_pos):
+            hidden_states[i][a + 1 : b] = images[idx]
+        return hidden_states.view(-1, 1, self.hidden_size)
+
+# llama2
+class LLAMA2Block(torch.nn.Module):
+    def __init__(self, block, block_id, hidden_size, final_layernorm = None):
+        super().__init__()
+        self.block = block
+        self.block_id = block_id
+        self.final_layernorm = final_layernorm
+        self.hidden_size = hidden_size
+
+    def forward(self, hidden_states, attention_mask, position_ids, past_kv):
+        hidden_states = hidden_states.view(1, -1, self.hidden_size)
+        hidden_states, presents = self.block(hidden_states,
+                                             attention_mask,
+                                             position_ids,
+                                             past_kv,
+                                             use_cache=True)
+        if self.final_layernorm is not None:
+            hidden_states = self.final_layernorm(hidden_states)
+            hidden_states = hidden_states.view(-1, self.hidden_size)[-1].view(1, 1, self.hidden_size)
+        if isinstance(presents, tuple):
+            presents = torch.stack(presents)
+        return hidden_states, presents
+
+class Llama2_7b_Chat(LLM):
+    def __init__(self, args):
+        self.model_name = 'Llama2_7b'
+        if 'Baichuan2' in args.path:
+            self.model_name = 'Baichuan2_7B'
+        if 'internlm' in args.path:
+            self.model_name = 'Internlm_7b'
+        if 'TinyLlama' in args.path:
+            self.model_name = 'TinyLlama'
+        if 'Yi' in args.path:
+            self.model_name = 'Yi'
+        if 'deepseek' in args.path:
+            self.model_name = 'deepseek'
+        if 'Llama-3' in args.path:
+            self.model_name = 'Llama3_8B'
+        super().__init__(args)
+
+    def load_model(self):
+        self.config = self.model.config
+        transformer = self.model.model
+        self.lm_ = self.model.lm_head
+        self.embed_ = transformer.embed_tokens
+        self.blocks_ = transformer.layers
+        self.final_layernorm_ = transformer.norm
+        # some wrapper
+        self.hidden_size = self.embed_.weight.shape[-1]
+        self.stop_id = self.tokenizer.eos_token_id
+        if hasattr(model, 'generation_config'):
+            self.stop_ids.append(self.stop_id)
+            self.stop_ids.append(self.model.generation_config.eos_token_id)
+        if self.model_name == 'Llama3_8B':
+            self.stop_ids.append(self.tokenizer.convert_tokens_to_ids("<|eot_id|>"))
+        self.block_nums = len(self.blocks_)
+        self.embed = Embedding(self.embed_, self.embed_bf16)
+        self.lm = Lm(self.lm_)
+        self.blocks = [LLAMA2Block(self.blocks_[i], i, self.hidden_size, self.final_layernorm_ if i == len(self.blocks_) - 1 else None) for i in range(self.block_nums)]
+        self.block_nums = self.config.num_hidden_layers
+        self.hidden_size = self.config.hidden_size
+        self.num_attention_heads = self.config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_attention_heads
+        self.num_key_value_heads = self.config.num_key_value_heads
+        self.past_kv_shape = [self.block_nums, 2, 1, self.num_key_value_heads, 0, self.head_dim]
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 3: "history_len" }
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 4: "history_len" }
+        }
+
+    def build_prompt(self, query):
+        if 'Baichuan2' in self.model_name:
+            return f'<reserved_106>{query}<reserved_107>'
+        if 'Internlm_7b' in self.model_name:
+            return f'<|User|>:{query}<eoh>\n<|Bot|>:'
+        if 'TinyLlama' in self.model_name:
+            return f'<s><|system|>\nYou are a friendly chatbot who always responds in the style of a pirate</s>\n<|user|>\n{query}</s>\n<|assistant|>\n'
+        if 'Yi' in self.model_name:
+            return f'<|im_start|> user\n{query}<|im_end|>\n<|im_start|> assistant\n'
+        if 'deepseek' in self.model_name:
+            return f'<|begin▁of▁sentence|>User: {query}\nAssistant:'
+        if 'Llama3' in self.model_name:
+            return f'<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{query}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n'
+        return f'[INST]{query}[/INST]'
+
+    def get_attention_mask(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.zeros([1, 1, 1, self.seq_len], dtype=torch.float32)
+        return (1 - torch.tril(torch.ones([1, 1, self.seq_len, self.seq_len]))) * torch.finfo(torch.float32).min
+
+    def get_position_ids(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.tensor([[self.seq_len - 1]], dtype=torch.long)
+        return torch.arange(self.seq_len, dtype=torch.long).unsqueeze(0)
+
+# phi-2
+class PHI2Block(torch.nn.Module):
+    def __init__(self, block, block_id, hidden_size):
+        super().__init__()
+        self.block = block
+        self.block_id = block_id
+        self.hidden_size = hidden_size
+
+    def forward(self, hidden_states, attention_mask, position_ids, past_kv):
+        theta = 1.0 / (10000 ** (torch.arange(0, 32, 2, dtype=torch.float32) / 32))
+        position_ids = position_ids.float().reshape(-1, 1)
+        idx_theta = position_ids * theta
+        rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=0).contiguous()
+        hidden_states = hidden_states.view(1, -1, self.hidden_size)
+        hidden_states, presents = self.block(hidden_states,
+                                             past_kv,
+                                             rotary_pos_emb=rotary_pos_emb,
+                                             causal_mask=attention_mask
+                                             )
+        if self.block_id == 31:
+            hidden_states = hidden_states[:, -1, :]
+        return hidden_states, presents
+
+class phi_2(LLM):
+    def __init__(self, args):
+        super().__init__(args)
+        self.model_name = 'phi-2'
+        self.asymmetric = False # TODO: some precision bug when using asymmetric
+
+    def load_model(self):
+        transformer = self.model.transformer
+        self.lm_ = self.model.lm_head
+        self.embed_ = transformer.embd.wte
+        self.hidden_size = self.embed_.weight.shape[-1]
+        self.blocks_ = transformer.h
+        # self.final_layernorm_ = transformer.final_layernorm
+        # some wrapper
+        self.stop_id = self.tokenizer.eos_token_id
+        self.block_nums = len(self.blocks_)
+        self.embed = Embedding(self.embed_, self.embed_bf16)
+        self.lm = Lm(self.lm_)
+        self.blocks = [PHI2Block(self.blocks_[i], i, self.hidden_size) for i in range(self.block_nums)]
+        # some config for export
+        self.past_kv_shape = [len(self.blocks), 1, 0, 2, 32, 80]
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 1: "history_len" }
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+            "past_key_values" : { 2: "history_len" }
+        }
+
+    def build_prompt(self, query):
+            return f'Instruct: {query}\nOutput:'
+
+    def get_attention_mask(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.zeros([1, 1, 1, 1]).bool()
+        attention_mask = ~torch.tril(torch.ones([1, 1, self.seq_len, self.seq_len]).bool())
+        return attention_mask
+
+    def get_position_ids(self) -> torch.Tensor:
+        if self.token_len:
+            return torch.tensor([[self.seq_len - 1]], dtype=torch.long)
+        return torch.arange(self.seq_len, dtype=torch.long).unsqueeze(0)
+
+# BGE is Embedding Model based Bert
+class BGEBlock(torch.nn.Module):
+    def __init__(self, block, block_id, hidden_size):
+        super().__init__()
+        self.block = block
+        self.block_id = block_id
+        self.hidden_size = hidden_size
+
+    def forward(self, hidden_states, attention_mask):
+        hidden_states = self.block(hidden_states, attention_mask)[0]
+        return hidden_states
+
+class bge(LLM):
+    def __init__(self, args):
+        super().__init__(args)
+        self.model_name = 'bge-large-zh'
+
+    def forward(self, input_ids, position_ids, attention_mask):
+        input_ids = input_ids.view(1, -1)
+        token_type_ids = (1 - attention_mask).view(1, -1)
+        hidden_states = self.embed(input_ids, token_type_ids, position_ids)[0].unsqueeze(0)
+        for i in range(self.block_nums):
+            hidden_states = self.blocks[i](hidden_states, attention_mask)
+        # hidden_states = self.lm(hidden_states) # sentence_embeddings not need
+        sentence_embeddings = hidden_states[:, 0]
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings
+
+    def response(self, query):
+        self.eval()
+        input_ids = self.tokenizer(query)['input_ids']
+        self.seq_len = len(input_ids)
+        input_ids = torch.tensor(input_ids)
+        position_ids = self.get_position_ids()
+        attention_mask = self.get_attention_mask()
+        res = self.forward(input_ids, position_ids, attention_mask)
+        return res
+
+    def load_model(self):
+        transformer = self.model.encoder
+        self.lm_ = self.model.pooler
+        self.embed_ = self.model.embeddings
+        self.hidden_size = self.embed_.word_embeddings.weight.shape[-1]
+        self.blocks_ = transformer.layer
+        # some wrapper
+        self.stop_id = self.tokenizer.eos_token_id
+        self.block_nums = len(self.blocks_)
+        self.embed = self.embed_
+        self.lm = self.lm_
+        self.blocks = [BGEBlock(self.blocks_[i], i, self.hidden_size) for i in range(self.block_nums)]
+        # some config for export
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "position_ids" : { 1: "seq_len" },
+            "attention_mask" : { 3: "seq_len" }
+        }
+
+    def export(self):
+        model = self.eval()
+        self.seq_len = 3
+        input_ids = torch.arange(3, dtype=torch.long)
+        position_ids = self.get_position_ids()
+        attention_mask = self.get_attention_mask()
+        onnx_model = f'./{self.onnx_path}/bge.onnx'
+        torch.onnx.export(
+            model, (input_ids, position_ids, attention_mask),
+            onnx_model,
+            verbose=self.export_verbose,
+            input_names=[
+                'input_ids',
+                'position_ids',
+                'attention_mask'
+            ],
+            output_names=['sentence_embeddings'],
+            dynamic_axes=self.model_dynamic_axes,
+            do_constant_folding=True,
+            opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+        if self.export_test:
+            self.seq_len = 4
+            position_ids = self.get_position_ids()
+            input_ids = torch.tensor([ 101,  872, 1962,  102 ], dtype=torch.long)
+            attention_mask = self.get_attention_mask()
+            # test
+            original_outs = model(input_ids, position_ids, attention_mask)
+            ort_session = ort.InferenceSession(onnx_model, providers=['CPUExecutionProvider'])
+            inputs = {
+                'input_ids' : input_ids.detach().numpy(),
+                'position_ids' : position_ids.detach().numpy(),
+                'attention_mask' : attention_mask.detach().numpy()
+            }
+            onnx_outs = ort_session.run(None, inputs)[0]
+            self.assert_equal(original_outs, onnx_outs)
+
+        token_str = None
+        if False: # save tokenizer in mnn
+            self.export_tokenizer()
+            token_path = os.path.join(self.onnx_path, "tokenizer.txt")
+            token_str = open(token_path, 'rt').read()
+
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path, 8, True, bizCode=token_str)
+
+    def get_position_ids(self) -> torch.Tensor:
+        return torch.arange(self.seq_len, dtype=torch.long).unsqueeze(0)
+
+    def get_attention_mask(self) -> torch.Tensor:
+        return torch.ones([1, 1, 1, self.seq_len], dtype=torch.long)
+
+class LoraModule(torch.nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.onnx_path = args.onnx_path
+        self.mnn_path = args.mnn_path
+        self.export_mnn = args.export_mnn
+        import peft
+        lora_weight = peft.load_peft_weights(args.path)
+        for k, v in lora_weight.items():
+            k = k.replace('.', '/')
+            self.register_buffer(k, v.cpu())
+
+    def forward(self, dummpy):
+        return self._buffers
+
+    def export(self):
+        onnx_model = f'./{self.onnx_path}/lora.onnx'
+        torch.onnx.export(self.eval(), torch.tensor([]), onnx_model)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
+
+class GOT(Qwen2_Chat):
+    def __init__(self, args):
+        super().__init__(args)
+    def load_hf(self, model_path: str):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model = GOTQwenForCausalLM.from_pretrained(model_path, trust_remote_code=True).float().eval()
+        self.config = self.model.config
+        if self.lora_path is not None:
+            adapter = PeftModel.from_pretrained(self.model, model_id=self.lora_path)
+            self.model = adapter.merge_and_unload(progressbar=True)
+    def load_model(self):
+        # Qwen2 models
+        self.model_name = 'GOT'
+        transformer = self.model.model
+        self.lm_ = self.model.lm_head
+        self.embed_ = transformer.embed_tokens
+        self.blocks_ = transformer.layers
+        self.final_layernorm_ = transformer.norm
+        self.visual = transformer.vision_tower_high
+        self.mm_projector_vary = transformer.mm_projector_vary
+        # some wrapper
+        self.stop_id = self.tokenizer.eos_token_id
+        if hasattr(self.model, 'generation_config'):
+            #self.stop_ids.append(self.stop_id)
+            #for id in self.model.generation_config.eos_token_id:
+            self.stop_ids.append(self.model.generation_config.eos_token_id)
+        self.block_nums = self.config.num_hidden_layers
+        self.hidden_size = self.config.hidden_size
+        self.image_size = self.hidden_size
+        self.image_token_len = self.config.image_token_len
+        self.num_heads = self.config.num_attention_heads
+        self.num_key_value_heads = self.config.num_key_value_heads
+        self.rope_theta = self.config.rope_theta
+        self.head_dim = self.hidden_size // self.num_heads
+        if self.embed_.weight is self.lm_.weight:
+            import copy
+            embed_copy = copy.deepcopy(self.embed_)
+            self.embed = GOTEmbedding(embed_copy, self.embed_bf16)
+        else:
+            self.embed = GOTEmbedding(self.embed_, self.embed_bf16)
+        self.lm = Lm(self.lm_)
+        self.past_kv_shape = [self.block_nums, 2, 1, 0, self.num_heads, self.head_dim]
+        #self.blocks = [QWEN2Block(self.model_name, self.blocks_[i], i, self.config, None) for i in range(self.block_nums)]
+        self.blocks = [Qwen2DecoderLayer(self.config,self.blocks_[i], i) for i in range(self.block_nums)]
+        #self.blocks = self.blocks_
+        # some config for export
+        self.block_dynamic_axes = {
+            "inputs_embeds" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+        }
+        self.model_dynamic_axes = {
+            "input_ids" : { 0: "seq_len" },
+            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
+            "position_ids" : { 0: "seq_len" },
+        }
+    def export_lm(self):
+        model = self.lm
+        hidden_states = torch.randn(1, self.hidden_size)
+        onnx_model = f'./{self.onnx_path}/lm.onnx'
+        torch.onnx.export(model, (hidden_states),
+                        onnx_model,
+                        verbose=self.export_verbose,
+                        input_names=['hidden_states'],
+                        output_names=['token_id'],
+                        do_constant_folding=True,
+                        dynamic_axes={
+                                    "hidden_states" : { 0: "seq_len" }
+                                },
+                        opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+    def export_norm(self):
+        model = self.final_layernorm_
+        hidden_states = torch.randn(1, self.image_token_len, self.hidden_size)
+        onnx_model = f'./{self.onnx_path}/norm.onnx'
+        torch.onnx.export(model, (hidden_states),
+                        onnx_model,
+                        verbose=self.export_verbose,
+                        input_names=['hidden_in'],
+                        output_names=['hidden_out'],
+                        do_constant_folding=True,
+                        dynamic_axes={
+                                    "hidden_in" : { 1: "seq_len" },
+                                    "hidden_out" : { 1: "seq_len"},
+                                },
+                        opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+    def export_projector_vary(self):
+        model = self.mm_projector_vary
+        hidden_states = torch.randn(1, self.image_token_len, self.hidden_size)
+        onnx_model = f'./{self.onnx_path}/mm_projector_vary.onnx'
+        torch.onnx.export(model, (hidden_states),
+                        onnx_model,
+                        verbose=self.export_verbose,
+                        input_names=['cnn_features'],
+                        output_names=['img_features'],
+                        do_constant_folding=True,
+                        opset_version=15)
+        if not self.skip_slim:
+            slim(onnx_model, output_model=onnx_model)
+            
+if __name__ == '__main__':
+    llm_models = {
+        'chatglm-6b': Chatglm_6b,
+        'chatglm2-6b': Chatglm2_6b,
+        'chatglm3-6b': Chatglm3_6b,
+        'codegeex2-6b': Chatglm2_6b,
+        'Qwen-7B-Chat': Qwen_Chat,
+        'Qwen-1_8B-Chat': Qwen_Chat,
+        'Qwen-1_8B': Qwen_Chat,
+        'Qwen-VL-Chat': Qwen_Chat,
+        'Qwen1_5-0_5B-Chat': Qwen2_Chat,
+        'Qwen1_5-1_8B-Chat': Qwen2_Chat,
+        'Qwen1_5-4B-Chat': Qwen2_Chat,
+        'Qwen1_5-7B-Chat': Qwen2_Chat,
+        'Baichuan2-7B-Chat': Llama2_7b_Chat,
+        'Llama-2-7b-chat-ms': Llama2_7b_Chat,
+        'Llama-3-8B-Instruct': Llama2_7b_Chat,
+        'internlm-chat-7b': Llama2_7b_Chat,
+        'TinyLlama-1_1B-Chat': Llama2_7b_Chat,
+        'Yi-6B-Chat': Llama2_7b_Chat,
+        'deepseek-llm-7b-chat': Llama2_7b_Chat,
+        'phi-2': phi_2,
+        'bge-large-zh': bge,
+        'lora': LoraModule,
+        'GOT':GOT
+    }
+    parser = argparse.ArgumentParser(description='llm_exporter', formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument('--path', type=str, default=r'D:\LearningCodes\GithubRepo\shouxieAI\GOT-OCR2.0\GOT-OCR-2.0-master\GOT_weights',
+                        help='path(`str` or `os.PathLike`):\nCan be either:'
+                        '\n\t- A string, the *model id* of a pretrained model like `THUDM/chatglm-6b`. [TODO]'
+                        '\n\t- A path to a *directory* clone from repo like `../chatglm-6b`.')
+    parser.add_argument('--type', type=str, choices=llm_models.keys(), default="GOT",
+                        help='type(`str`, *optional*):'
+                        '\n\tThe pretrain llm model type.'
+                        )
+    parser.add_argument('--lora_path', type=str, default=None, help='lora path, defaut is `None` mean not apply lora.')
+    parser.add_argument('--onnx_path', type=str, default='./onnx', help='export onnx model path, defaut is `./onnx`.')
+    parser.add_argument('--mnn_path', type=str, default='./mnn', help='export mnn model path, defaut is `./mnn`.')
+    parser.add_argument('--export_mnn', action='store_true', default=False, help='Whether or not to export mnn model after onnx.')
+    parser.add_argument('--export_verbose', action='store_true', default=False, help='Whether or not to export onnx with verbose.')
+    parser.add_argument('--export_test', action='store_true', help='Whether or not to export onnx with test using onnxruntime.')
+    parser.add_argument('--test', type=str, help='test model inference with query `TEST`.')
+    parser.add_argument('--export', action='store_true', help='export model to an `onnx` model.')
+    parser.add_argument('--export_split', default=True,
+                        help='export model split to some `onnx` models:'
+                        '\n\t- embedding model.'
+                        '\n\t- block models.'
+                        '\n\t- lm_head model.'
+                        )
+    parser.add_argument('--export_token', action='store_true', help='export llm tokenizer to a txt file.')
+    parser.add_argument('--export_embed', action='store_true', help='export llm embedding to an `onnx` model.')
+    parser.add_argument('--export_visual', action='store_true', help='export llm visual model to an `onnx` model.')
+    parser.add_argument('--export_lm', action='store_true', help='export llm lm_head to an `onnx` model.')
+    parser.add_argument('--export_block', type=int, help='export llm block [id] to an `onnx` model.')
+    parser.add_argument('--export_blocks', action='store_true', help='export llm all blocks to `onnx` models.')
+    parser.add_argument('--embed_bin', action='store_true', help='export embedding weight as bin file with dtype `bfloat16`')
+    parser.add_argument('--embed_bf16', action='store_true', help='using `bfloat16` replace `float32` in embedding.')
+    parser.add_argument('--skip_slim', action='store_true', help='Whether or not to skip onnx-slim.')
+
+
+    args = parser.parse_args()
+    model_path = args.path
+    model_type = args.type
+    # not sepcify model type, using path
+    if model_type is None:
+        for model in llm_models:
+            if model in model_path:
+                model_type = model
+    if model_type is None:
+        raise RuntimeError('Please specify model type.')
+
+    # # copy modeling py file to pretrain model for export
+    # for file in glob.glob(f'./llm_models/{model_type}/*'):
+    #     shutil.copy2(file, model_path)
+
+    llm_exporter = llm_models[model_type](args)
+
+    # some actions
+    if args.test is not None:
+        llm_exporter.response(args.test)
+
+    if args.export:
+        llm_exporter.export()
+
+    if args.export_token:
+        llm_exporter.export_tokenizer()
+
+    if args.export_embed or args.export_split:
+        llm_exporter.export_embed()
+
+    if args.export_visual or args.export_split:
+        llm_exporter.export_visual()
+
+    if args.export_lm or args.export_split:
+        llm_exporter.export_lm()
+        llm_exporter.export_projector_vary()
+        llm_exporter.export_norm()
+
+    if args.export_blocks or args.export_split:
+        llm_exporter.export_blocks()
+
+    if args.export_block is not None:
+        llm_exporter.export_block(args.export_block)