Add onnx model

.gitattributes

@@ -14,6 +14,7 @@
 *.npy filter=lfs diff=lfs merge=lfs -text
 *.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
+*.onnx.data filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

bgem3_model.py

@@ -0,0 +1,80 @@
+# modified from https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/BGE_M3/modeling.py
+
+import os
+import torch
+from torch import nn, Tensor
+from transformers import AutoModel, AutoConfig
+from huggingface_hub import snapshot_download
+from typing import Dict
+
+
+class BGEM3InferenceModel(nn.Module):
+    def __init__(
+        self,
+        model_name: str = "BAAI/bge-m3",
+        colbert_dim: int = -1,
+    ) -> None:
+        super().__init__()
+
+        model_name = snapshot_download(
+            repo_id=model_name,
+            allow_patterns=[
+                "model.safetensors",
+                "colbert_linear.pt",
+                "sparse_linear.pt",
+                "config.json",
+            ],
+        )
+
+        self.config = AutoConfig.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+        self.colbert_linear = torch.nn.Linear(
+            in_features=self.model.config.hidden_size,
+            out_features=(
+                self.model.config.hidden_size if colbert_dim == -1 else colbert_dim
+            ),
+        )
+        self.sparse_linear = torch.nn.Linear(
+            in_features=self.model.config.hidden_size, out_features=1
+        )
+        colbert_state_dict = torch.load(
+            os.path.join(model_name, "colbert_linear.pt"), map_location="cpu"
+        )
+        sparse_state_dict = torch.load(
+            os.path.join(model_name, "sparse_linear.pt"), map_location="cpu"
+        )
+        self.colbert_linear.load_state_dict(colbert_state_dict)
+        self.sparse_linear.load_state_dict(sparse_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:])
+        colbert_vecs = colbert_vecs * attention_mask[:, 1:][:, :, None].float()
+        return colbert_vecs
+
+    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

config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": ".",
+  "architectures": [
+    "XLMRobertaModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 8194,
+  "model_type": "xlm-roberta",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.37.2",
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "vocab_size": 250002
+}

export_onnx.py

@@ -0,0 +1,85 @@
+import copy
+import argparse
+from optimum.exporters.onnx import onnx_export_from_model
+from collections import OrderedDict
+from typing import Dict
+from optimum.exporters.onnx.model_configs import XLMRobertaOnnxConfig
+from bgem3_model import BGEM3InferenceModel
+
+
+class BGEM3OnnxConfig(XLMRobertaOnnxConfig):
+    @property
+    def outputs(self) -> Dict[str, Dict[int, str]]:
+        """
+        Dict containing the axis definition of the output tensors to provide to the model.
+
+        Returns:
+            `Dict[str, Dict[int, str]]`: A mapping of each output name to a mapping of axis position to the axes symbolic name.
+        """
+        return copy.deepcopy(
+            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, opset: int, device: str, optimize: str, atol: str):
+    model = BGEM3InferenceModel()
+    bgem3_onnx_config = BGEM3OnnxConfig(model.config)
+    onnx_export_from_model(
+        model,
+        output=output,
+        task="feature-extraction",
+        custom_onnx_configs={"model": bgem3_onnx_config},
+        opset=opset,
+        optimize=optimize,
+        atol=atol,
+        device=device,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--output",
+        type=str,
+        help="Path indicating the directory where to store the generated ONNX model.",
+    )
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=None,
+        help="If specified, ONNX opset version to export the model with. Otherwise, the default opset for the given model architecture will be used.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help='The device to use to do the export. Defaults to "cpu".',
+    )
+    parser.add_argument(
+        "--optimize",
+        type=str,
+        default=None,
+        choices=["O1", "O2", "O3", "O4"],
+        help=(
+            "Allows to run ONNX Runtime optimizations directly during the export. Some of these optimizations are specific to ONNX Runtime, and the resulting ONNX will not be usable with other runtime as OpenVINO or TensorRT. Possible options:\n"
+            "    - O1: Basic general optimizations\n"
+            "    - O2: Basic and extended general optimizations, transformers-specific fusions\n"
+            "    - O3: Same as O2 with GELU approximation\n"
+            "    - O4: Same as O3 with mixed precision (fp16, GPU-only, requires `--device cuda`)"
+        ),
+    )
+    parser.add_argument(
+        "--atol",
+        type=float,
+        default=None,
+        help="If specified, the absolute difference tolerance when validating the model. Otherwise, the default atol for the model will be used.",
+    )
+    args = parser.parse_args()
+
+    main(args.output, args.opset, args.device, args.optimize, args.atol)

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:138d09ae2920b7e8731f01cba6b5ad996fd64bdfe34971e2d22ecbcf322e25b1
+size 108605

model.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfa52ffdb65db76612d6c3ad92130221822f613004113e8c0af18c5eab81a81d
+size 2271088656

ort_config.json

@@ -0,0 +1,39 @@
+{
+  "one_external_file": true,
+  "opset": null,
+  "optimization": {
+    "disable_attention": null,
+    "disable_attention_fusion": false,
+    "disable_bias_gelu": null,
+    "disable_bias_gelu_fusion": false,
+    "disable_bias_skip_layer_norm": null,
+    "disable_bias_skip_layer_norm_fusion": false,
+    "disable_embed_layer_norm": true,
+    "disable_embed_layer_norm_fusion": true,
+    "disable_gelu": null,
+    "disable_gelu_fusion": false,
+    "disable_group_norm_fusion": true,
+    "disable_layer_norm": null,
+    "disable_layer_norm_fusion": false,
+    "disable_packed_kv": true,
+    "disable_rotary_embeddings": false,
+    "disable_shape_inference": true,
+    "disable_skip_layer_norm": null,
+    "disable_skip_layer_norm_fusion": false,
+    "enable_gelu_approximation": false,
+    "enable_gemm_fast_gelu_fusion": false,
+    "enable_transformers_specific_optimizations": true,
+    "fp16": false,
+    "no_attention_mask": false,
+    "optimization_level": 2,
+    "optimize_for_gpu": false,
+    "optimize_with_onnxruntime_only": null,
+    "use_mask_index": false,
+    "use_multi_head_attention": false,
+    "use_raw_attention_mask": false
+  },
+  "optimum_version": "1.17.0",
+  "quantization": {},
+  "transformers_version": "4.37.2",
+  "use_external_data_format": true
+}

requirements.txt

@@ -0,0 +1,7 @@
+accelerate==0.27.2
+huggingface-hub==0.20.3
+onnx==1.15.0
+onnxruntime==1.17.0
+optimum==1.17.0
+torch==2.2.0
+transformers==4.37.2

sentencepiece.bpe.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865
+size 5069051

special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6710678b12670bc442b99edc952c4d996ae309a7020c1fa0096dd245c2faf790
+size 17082821

tokenizer_config.json

@@ -0,0 +1,55 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "250001": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_max_length": 8192,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}