|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import json |
|
|
|
import numpy as np |
|
import triton_python_backend_utils as pb_utils |
|
from transformers import AutoTokenizer, LlamaTokenizer, T5Tokenizer |
|
|
|
|
|
class TritonPythonModel: |
|
"""Your Python model must use the same class name. Every Python model |
|
that is created must have "TritonPythonModel" as the class name. |
|
""" |
|
|
|
def initialize(self, args): |
|
"""`initialize` is called only once when the model is being loaded. |
|
Implementing `initialize` function is optional. This function allows |
|
the model to initialize any state associated with this model. |
|
Parameters |
|
---------- |
|
args : dict |
|
Both keys and values are strings. The dictionary keys and values are: |
|
* model_config: A JSON string containing the model configuration |
|
* model_instance_kind: A string containing model instance kind |
|
* model_instance_device_id: A string containing model instance device ID |
|
* model_repository: Model repository path |
|
* model_version: Model version |
|
* model_name: Model name |
|
""" |
|
|
|
model_config = json.loads(args['model_config']) |
|
tokenizer_dir = model_config['parameters']['tokenizer_dir'][ |
|
'string_value'] |
|
tokenizer_type = model_config['parameters']['tokenizer_type'][ |
|
'string_value'] |
|
|
|
if tokenizer_type == 't5': |
|
self.tokenizer = T5Tokenizer(vocab_file=tokenizer_dir, |
|
padding_side='left') |
|
elif tokenizer_type == 'auto': |
|
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir, |
|
padding_side='left') |
|
elif tokenizer_type == 'llama': |
|
self.tokenizer = LlamaTokenizer.from_pretrained( |
|
tokenizer_dir, legacy=False, padding_side='left') |
|
else: |
|
raise AttributeError( |
|
f'Unexpected tokenizer type: {tokenizer_type}') |
|
self.tokenizer.pad_token = self.tokenizer.eos_token |
|
|
|
self._init_token_map() |
|
|
|
|
|
output_config = pb_utils.get_output_config_by_name( |
|
model_config, "OUTPUT") |
|
|
|
|
|
self.output_dtype = pb_utils.triton_string_to_numpy( |
|
output_config['data_type']) |
|
output_lens_config = pb_utils.get_output_config_by_name( |
|
model_config, "OUTPUT_LENS") |
|
|
|
|
|
self.output_lens_dtype = pb_utils.triton_string_to_numpy( |
|
output_lens_config['data_type']) |
|
|
|
def _init_token_map(self): |
|
v = self.tokenizer.get_vocab() |
|
self.token_map = [None] * len(v) |
|
for k, val in v.items(): |
|
self.token_map[val] = k |
|
|
|
for i in range(len(v)): |
|
if self.token_map[i] is None: |
|
print("error %s" % i) |
|
|
|
def execute(self, requests): |
|
"""`execute` must be implemented in every Python model. `execute` |
|
function receives a list of pb_utils.InferenceRequest as the only |
|
argument. This function is called when an inference is requested |
|
for this model. Depending on the batching configuration (e.g. Dynamic |
|
Batching) used, `requests` may contain multiple requests. Every |
|
Python model, must create one pb_utils.InferenceResponse for every |
|
pb_utils.InferenceRequest in `requests`. If there is an error, you can |
|
set the error argument when creating a pb_utils.InferenceResponse. |
|
Parameters |
|
---------- |
|
requests : list |
|
A list of pb_utils.InferenceRequest |
|
Returns |
|
------- |
|
list |
|
A list of pb_utils.InferenceResponse. The length of this list must |
|
be the same as `requests` |
|
""" |
|
|
|
responses = [] |
|
|
|
|
|
|
|
for idx, request in enumerate(requests): |
|
|
|
tokens_batch = pb_utils.get_input_tensor_by_name( |
|
request, 'TOKENS_BATCH').as_numpy() |
|
|
|
|
|
|
|
|
|
|
|
|
|
outputs, output_lens = self._postprocessing(tokens_batch) |
|
|
|
|
|
|
|
output_tensor = pb_utils.Tensor( |
|
'OUTPUT', |
|
np.array(outputs).astype(self.output_dtype)) |
|
output_lens_tensor = pb_utils.Tensor( |
|
'OUTPUT_LENS', |
|
np.array(output_lens).astype(self.output_lens_dtype)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
inference_response = pb_utils.InferenceResponse( |
|
output_tensors=[output_tensor, output_lens_tensor]) |
|
responses.append(inference_response) |
|
|
|
|
|
|
|
return responses |
|
|
|
def finalize(self): |
|
"""`finalize` is called only once when the model is being unloaded. |
|
Implementing `finalize` function is optional. This function allows |
|
the model to perform any necessary clean ups before exit. |
|
""" |
|
print('Cleaning up...') |
|
|
|
def _single_token_decode(self, token): |
|
st = self.token_map[token] |
|
if st[0] == '▁': |
|
return " " + st[1:] |
|
return st |
|
|
|
def _postprocessing(self, tokens_batch): |
|
outputs = [] |
|
output_lens = [] |
|
for beam_tokens in tokens_batch: |
|
total_len = 0 |
|
for tokens in beam_tokens: |
|
if len(tokens) == 1: |
|
output = self._single_token_decode(tokens[0]) |
|
else: |
|
output = self.tokenizer.decode(tokens) |
|
print(output) |
|
outputs.append(output.encode('utf8')) |
|
total_len += len(tokens) |
|
output_lens.append(total_len) |
|
return outputs, output_lens |
|
|