infgrad

/

puff-base-v1

like 9

不要用optimum,不好用，自己手动转一下，参考代码：

import torch
import torch.nn as nn
from os.path import join
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel


class VECModelWrapper(nn.Module):
    def __init__(self, model_dir: str):
        super().__init__()
        self.model = AutoModel.from_pretrained(model_dir).cpu()
        t = AutoTokenizer.from_pretrained(model_dir)
        print(t([["123", "4,5,6"], ["123", "4,5,6"]]).keys())
        # for k,v in self.model.named_parameters():
        #     print(k,v.shape)

        self.vector_linear = torch.nn.Linear(in_features=768, out_features=4096)
        vector_linear_dict = {
            k.replace("linear.", ""): v for k, v in
            torch.load(join(model_dir, f"2_Dense_4096/pytorch_model.bin"), map_location="cpu").items()
        }
        self.vector_linear.load_state_dict(vector_linear_dict)
        self.model.eval()
        self.vector_linear.eval()

    def forward(
            self,
            input_ids,
            # token_type_ids,
            attention_mask
    ):
        last_hidden_state = self.model(
            input_ids=input_ids,
            # token_type_ids=token_type_ids,
            attention_mask=attention_mask
        )[0]
        last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
        vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
        vectors = self.vector_linear(vectors)
        print(vectors.shape)
        print(vectors[:, :4])
        return vectors


def convert_pt_to_onnx(model_dir="./output/tiny_bert/best-f1", save_path="sim_model.onnx"):
    # 加载模型
    model = VECModelWrapper(model_dir=model_dir)

    # 转为onnx
    # 要准备一个输入以便跟踪模型结构，随意一个输入就行
    input_ids = torch.tensor([[1, 2, 3], [11, 12, 113]], dtype=torch.long)
    attention_mask = torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.long)
    # token_type_ids = torch.tensor([[0, 0, 0], [0, 1, 0]], dtype=torch.long)
    torch.onnx.export(model,  # 模型
                      args=(
                          input_ids,
                          attention_mask,
                          # token_type_ids
                      ),  # 模型输入
                      f=save_path,  # 存储路径
                      export_params=True,
                      opset_version=14,  # op版本，越高越好
                      do_constant_folding=True,
                      input_names=[
                          'input_ids',
                          "attention_mask",
                          # "token_type_ids"
                      ],  # 模型输入的名字，随便写
                      output_names=['vectors'],  # 模型输出的名字，使用ort就要根据这个名字获取结果
                      dynamic_axes={"input_ids": [0, 1],  # 动态维度,第0和1维度都是可变的
                                    "attention_mask": [0, 1],
                                    # "token_type_ids": [0, 1],
                                    "vectors": [0],
                                    })


if __name__ == "__main__":
    model_dir = r"E:\pt\puff-base-v1"
    # model_dir = r"E:\PublicModels\bge-small-zh-v1.5"
    onnx_path = r"E:\onnx\puff-base-v1\1\model.onnx"

    convert_pt_to_onnx(model_dir=model_dir, save_path=onnx_path)

谢谢！