Upload convert.py

convert.py

@@ -0,0 +1,130 @@
+from __future__ import annotations
+
+import os
+from collections import OrderedDict
+from pathlib import Path
+from typing import Dict
+
+import torch
+from huggingface_hub import snapshot_download
+from optimum.exporters.onnx import export
+from optimum.exporters.onnx.model_configs import XLMRobertaOnnxConfig
+from optimum.onnxruntime import ORTModelForCustomTasks, ORTOptimizer
+from optimum.onnxruntime.configuration import AutoOptimizationConfig
+from torch import Tensor
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel, XLMRobertaConfig
+
+
+class BGEM3InferenceModel(PreTrainedModel):
+    config_class = XLMRobertaConfig
+    base_model_prefix = "BGEM3InferenceModel"
+    model_tags = ["BAAI/bge-m3"]
+
+    def __init__(self, model_name: str = "BAAI/bge-m3"):
+        super().__init__(PretrainedConfig())
+
+        model_name = snapshot_download(repo_id=model_name)
+
+        self.config = AutoConfig.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+
+        self.sparse_linear = torch.nn.Linear(
+            in_features=self.model.config.hidden_size,
+            out_features=1,
+        )
+        sparse_state_dict = torch.load(os.path.join(model_name, "sparse_linear.pt"), map_location="cpu")
+        self.sparse_linear.load_state_dict(sparse_state_dict)
+
+        self.colbert_linear = torch.nn.Linear(
+            in_features=self.model.config.hidden_size,
+            out_features=self.model.config.hidden_size,
+        )
+        colbert_state_dict = torch.load(os.path.join(model_name, "colbert_linear.pt"), map_location="cpu")
+        self.colbert_linear.load_state_dict(colbert_state_dict)
+
+    def dense_embedding(self, last_hidden_state: Tensor) -> Tensor:
+        return last_hidden_state[:, 0]
+
+    def sparse_embedding(self, last_hidden_state: Tensor) -> Tensor:
+        with torch.no_grad():
+            return torch.relu(self.sparse_linear(last_hidden_state))
+
+    def colbert_embedding(self, last_hidden_state: Tensor, attention_mask: Tensor) -> Tensor:
+        with torch.no_grad():
+            colbert_vecs = self.colbert_linear(last_hidden_state[:, 1:])
+        return colbert_vecs * attention_mask[:, 1:][:, :, None].float()
+
+    def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Dict[str, Tensor]:
+        with torch.no_grad():
+            last_hidden_state = self.model(
+                input_ids=input_ids, attention_mask=attention_mask, return_dict=True
+            ).last_hidden_state
+
+        output = {}
+        dense_vecs = self.dense_embedding(last_hidden_state)
+        output["dense_vecs"] = torch.nn.functional.normalize(dense_vecs, dim=-1)
+
+        sparse_vecs = self.sparse_embedding(last_hidden_state)
+        output["sparse_vecs"] = sparse_vecs
+
+        colbert_vecs = self.colbert_embedding(last_hidden_state, attention_mask)
+        output["colbert_vecs"] = torch.nn.functional.normalize(colbert_vecs, dim=-1)
+
+        return output
+
+
+class BGEM3OnnxConfig(XLMRobertaOnnxConfig):
+    @property
+    def outputs(self) -> Dict[str, Dict[int, str]]:
+        return OrderedDict(
+            {
+                "dense_vecs": {0: "batch_size", 1: "embedding"},
+                "sparse_vecs": {0: "batch_size", 1: "token", 2: "weight"},
+                "colbert_vecs": {0: "batch_size", 1: "token", 2: "embedding"},
+            }
+        )
+
+
+def main(output: str, device: str = "cuda", optimize: str = "O4"):
+    # 加载模型
+    model = BGEM3InferenceModel()
+    model.save_pretrained(output)
+
+    # 配置
+    bgem3_onnx_config = BGEM3OnnxConfig(model.config)
+
+    # 导出
+    export(
+        model,
+        output=Path(output) / "model.onnx",
+        config=bgem3_onnx_config,
+        opset=bgem3_onnx_config.DEFAULT_ONNX_OPSET,
+        device=device,
+    )
+
+    optimizer = ORTOptimizer.from_pretrained(output, file_names=["model.onnx"])
+    optimization_config = AutoOptimizationConfig.with_optimization_level(optimization_level=optimize)
+    optimization_config.disable_shape_inference = True
+    if optimize == "O4":
+        optimization_config.optimize_for_gpu = True
+        optimization_config.fp16 = True
+        optimization_config.optimization_level = 99
+    optimizer.optimize(save_dir=output, optimization_config=optimization_config, file_suffix="")
+
+    ORTModelForCustomTasks.from_pretrained(
+        output,
+        provider="CUDAExecutionProvider" if device == "cuda" else "CPUExecutionProvider",
+    )
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output", type=str)
+    parser.add_argument("--device", type=str, choices=["cuda", "cpu"], default="cuda")
+    parser.add_argument("--optimize", type=str, choices=["O1", "O2", "O3", "O4"], default="O4")
+    parser.add_argument("--push_to_hub", action="store_true", default=False)
+    parser.add_argument("--push_to_hub_repo_id", type=str, default="JeremyHibiki/bge-m3-onnx")
+    args = parser.parse_args()
+    main(args.output, args.device, args.optimize)