chore: add onnx model

viencoder.model/1/model.onnx

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

viencoder.model/config.pbtxt

@@ -0,0 +1,24 @@
+name: "viencoder.model"
+platform: "onnxruntime_onnx"
+backend: "onnxruntime"
+default_model_filename: "model.onnx"
+max_batch_size: 0
+input [
+  {
+    name: "input_ids"
+    data_type: TYPE_INT64
+    dims: [ -1, -1 ]
+  },
+  {
+    name: "attention_mask"
+    data_type: TYPE_INT64
+    dims: [ -1, -1 ]
+  }
+]
+output [
+  {
+    name: "sentence_embedding"
+    data_type: TYPE_FP32
+    dims: [ -1, 768 ]
+  }
+]

viencoder.tokenizer/1/added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<mask>": 64000
+}

viencoder.tokenizer/1/model.py

@@ -0,0 +1,51 @@
+
+import os
+from typing import Dict, List
+
+import numpy as np
+import triton_python_backend_utils as pb_utils
+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(
+              query, padding=True, truncation=True,return_tensors=TensorType.NUMPY
+          )
+          # tensorrt uses int32 as input type, ort uses int64
+          tokens = {k: v.astype(np.int64) 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

viencoder.tokenizer/1/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": false,
+    "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
+  }
+}

viencoder.tokenizer/1/tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "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
+    },
+    "64000": {
+      "content": "<mask>",
+      "lstrip": false,
+      "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": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "PhobertTokenizer",
+  "unk_token": "<unk>"
+}

viencoder.tokenizer/config.pbtxt

@@ -0,0 +1,24 @@
+name: "viencoder.tokenizer"
+max_batch_size: 0
+backend: "python"
+
+input [
+{
+  name: "TEXT"
+  data_type: TYPE_STRING
+  dims: [ -1 ]
+}
+]
+
+output [
+{
+  name: "input_ids"
+  data_type: TYPE_INT64
+  dims: [-1, -1]
+},
+{
+  name: "attention_mask"
+  data_type: TYPE_INT64
+  dims: [-1, -1]
+}
+]

viencoder/config.pbtxt

@@ -0,0 +1,64 @@
+name: "viencoder"
+# maximum batch size 
+max_batch_size: 0 
+platform: "ensemble"
+
+#input to the model 
+input [
+{
+  name: "TEXT"
+  data_type: TYPE_STRING
+  dims: [ -1 ] 
+  # -1 means dynamic axis, aka this dimension may change 
+}
+]
+
+#output of the model 
+output {
+  name: "output_0"
+  data_type: TYPE_FP32
+  dims: [-1, -1] 
+  # two dimensional tensor, where 1st dimension: batch-size, 2nd dimension: #classes
+}
+
+#Type of scheduler to be used
+ensemble_scheduling {
+  step [
+    {
+      model_name: "viencoder.tokenizer"
+      model_version: -1
+      input_map {
+      key: "TEXT"
+      value: "TEXT"
+    }
+    output_map [
+    {
+      key: "input_ids"
+      value: "input_ids"
+    },
+    {
+      key: "attention_mask"
+      value: "attention_mask"
+    }
+    ]
+    },
+    {
+      model_name: "viencoder.model"
+      model_version: -1
+    input_map [
+      {
+        key: "input_ids"
+        value: "input_ids"
+      },
+      {
+        key: "attention_mask"
+        value: "attention_mask"
+      }
+    ]
+    output_map {
+        key: "sentence_embedding"
+        value: "output_0"
+      }
+    }
+  ]
+}