feat: ONNX model

Signed-off-by: Thytu <valentin.de-matos@epitech.eu>

sentence-transformers_all-MiniLM-L6-v2_onnx_inference/config.pbtxt

@@ -0,0 +1,66 @@
+name: "sentence-transformers_all-MiniLM-L6-v2_onnx_inference"
+max_batch_size: 0
+platform: "ensemble"
+
+input [
+{
+    name: "TEXT"
+    data_type: TYPE_STRING
+    dims: [ -1 ]
+}
+]
+
+output {
+    name: "output"
+    data_type: TYPE_FP32
+    dims: [-1, 384]
+}
+
+ensemble_scheduling {
+    step [
+        {
+            model_name: "sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize"
+            model_version: -1
+            input_map {
+            key: "TEXT"
+            value: "TEXT"
+        }
+        output_map [
+{
+    key: "input_ids"
+    value: "input_ids"
+},
+{
+    key: "token_type_ids"
+    value: "token_type_ids"
+},
+{
+    key: "attention_mask"
+    value: "attention_mask"
+}
+        ]
+        },
+        {
+            model_name: "sentence-transformers_all-MiniLM-L6-v2_onnx_model"
+            model_version: -1
+            input_map [
+{
+    key: "input_ids"
+    value: "input_ids"
+},
+{
+    key: "token_type_ids"
+    value: "token_type_ids"
+},
+{
+    key: "attention_mask"
+    value: "attention_mask"
+}
+            ]
+        output_map {
+                key: "output"
+                value: "output"
+            }
+        }
+    ]
+}

sentence-transformers_all-MiniLM-L6-v2_onnx_model/1/model.bin

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

sentence-transformers_all-MiniLM-L6-v2_onnx_model/config.pbtxt

@@ -0,0 +1,35 @@
+name: "sentence-transformers_all-MiniLM-L6-v2_onnx_model"
+max_batch_size: 0
+platform: "onnxruntime_onnx"
+default_model_filename: "model.bin"
+
+input [
+{
+    name: "input_ids"
+    data_type: TYPE_INT32
+    dims: [-1, -1]
+},
+{
+    name: "token_type_ids"
+    data_type: TYPE_INT32
+    dims: [-1, -1]
+},
+{
+    name: "attention_mask"
+    data_type: TYPE_INT32
+    dims: [-1, -1]
+}
+]
+
+output {
+    name: "output"
+    data_type: TYPE_FP32
+    dims: [-1, 384]
+}
+
+instance_group [
+    {
+      count: 1
+      kind: KIND_GPU
+    }
+]

sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize/1/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "sentence-transformers/all-MiniLM-L6-v2",
+  "architectures": [
+    "BertModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1536,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 6,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.24.0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize/1/model.py

@@ -0,0 +1,70 @@
+#  Copyright 2022, Lefebvre Dalloz Services
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This module is copy-pasted in generated Triton configuration folder to perform the tokenization step.
+"""
+
+# noinspection DuplicatedCode
+import os
+from typing import Dict, List
+
+import numpy as np
+
+
+try:
+    # noinspection PyUnresolvedReferences
+    import triton_python_backend_utils as pb_utils
+except ImportError:
+    pass  # triton_python_backend_utils exists only inside Triton Python backend.
+
+from transformers import AutoTokenizer, PreTrainedTokenizer, TensorType
+
+
+class TritonPythonModel:
+    tokenizer: PreTrainedTokenizer
+
+    def initialize(self, args: Dict[str, str]) -> None:
+        """
+        Initialize the tokenization process
+        :param args: arguments from Triton config file
+        """
+        # more variables in https://github.com/triton-inference-server/python_backend/blob/main/src/python.cc
+        path: str = os.path.join(args["model_repository"], args["model_version"])
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+
+    def execute(self, requests) -> "List[List[pb_utils.Tensor]]":
+        """
+        Parse and tokenize each request
+        :param requests: 1 or more requests received by Triton server.
+        :return: text as input tensors
+        """
+        responses = []
+        # for loop for batch requests (disabled in our case)
+        for request in requests:
+            # binary data typed back to string
+            query = [t.decode("UTF-8") for t in pb_utils.get_input_tensor_by_name(request, "TEXT").as_numpy().tolist()]
+            tokens: Dict[str, np.ndarray] = self.tokenizer(text=query, return_tensors=TensorType.NUMPY)
+            # tensorrt uses int32 as input type, ort uses int64
+            tokens = {k: v.astype(np.int32) for k, v in tokens.items()}
+            # communicate the tokenization results to Triton server
+            outputs = list()
+            for input_name in self.tokenizer.model_input_names:
+                tensor_input = pb_utils.Tensor(input_name, tokens[input_name])
+                outputs.append(tensor_input)
+
+            inference_response = pb_utils.InferenceResponse(output_tensors=outputs)
+            responses.append(inference_response)
+
+        return responses

sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize/1/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize/1/tokenizer_config.json

@@ -0,0 +1,16 @@
+{
+  "cls_token": "[CLS]",
+  "do_basic_tokenize": true,
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "model_max_length": 512,
+  "name_or_path": "sentence-transformers/all-MiniLM-L6-v2",
+  "never_split": null,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "special_tokens_map_file": "/root/.cache/huggingface/hub/models--sentence-transformers--all-MiniLM-L6-v2/snapshots/7dbbc90392e2f80f3d3c277d6e90027e55de9125/special_tokens_map.json",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "BertTokenizer",
+  "unk_token": "[UNK]"
+}

sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize/config.pbtxt

@@ -0,0 +1,36 @@
+name: "sentence-transformers_all-MiniLM-L6-v2_onnx_tokenize"
+max_batch_size: 0
+backend: "python"
+
+input [
+{
+    name: "TEXT"
+    data_type: TYPE_STRING
+    dims: [ -1 ]
+}
+]
+
+output [
+{
+    name: "input_ids"
+    data_type: TYPE_INT32
+    dims: [-1, -1]
+},
+{
+    name: "token_type_ids"
+    data_type: TYPE_INT32
+    dims: [-1, -1]
+},
+{
+    name: "attention_mask"
+    data_type: TYPE_INT32
+    dims: [-1, -1]
+}
+]
+
+instance_group [
+    {
+      count: 1
+      kind: KIND_GPU
+    }
+]