File size: 4,829 Bytes
966f777 a6635ad 966f777 a6635ad 966f777 a6635ad 966f777 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
"""Postprocessor for codegen-350M-mono-gptj."""
import json
from pathlib import Path
from typing import Any, Dict, List
import numpy as np
import triton_python_backend_utils as pb_utils
from transformers import AutoTokenizer
class TritonPythonModel:
"""Postprocessor for codegen-350M-mono-gptj."""
def initialize(self, args: Dict[str, Any]) -> None:
"""`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.
Args:
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
"""
# Parse model configs
self.model_config = model_config = json.loads(args["model_config"])
# Parse model output configs
output_config = pb_utils.get_output_config_by_name(model_config, "OUTPUT")
# Convert Triton types to numpy types
self.output_dtype = pb_utils.triton_string_to_numpy(output_config["data_type"])
# Init a tokenizer for postprocessing.
cur_folder = Path(__file__).parent
cache_dir = cur_folder / ".cache"
self.tokenizer = AutoTokenizer.from_pretrained(
"Salesforce/codegen-350M-mono", cache_dir=cache_dir
)
def execute(
self, requests: List["pb_utils.InferenceRequest"]
) -> List["pb_utils.InferenceResponse"]:
"""Preprocess the 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.
Args:
requests : A list of pb_utils.InferenceRequest
Returns:
A list of pb_utils.InferenceResponse. The length of this list must
be the same as `requests`
"""
responses = []
# Every Python backend must iterate over everyone of the requests
# and create a pb_utils.InferenceResponse for each of them.
for request in requests:
# Get input tensors
tokens_batch = pb_utils.get_input_tensor_by_name(
request, "TOKENS_BATCH"
).as_numpy()
# Reshape Input
# tokens_batch = tokens_batch.reshape([-1, tokens_batch.shape[0]])
# tokens_batch = tokens_batch.T
# Postprocessing output data.
outputs = self._postprocessing(tokens_batch)
# Create output tensors. You need pb_utils.Tensor
# objects to create pb_utils.InferenceResponse.
output_tensor = pb_utils.Tensor(
"OUTPUT", np.array(outputs).astype(self.output_dtype)
)
# Create InferenceResponse. You can set an error here in case
# there was a problem with handling this inference request.
# Below is an example of how you can set errors in inference
# response:
#
# pb_utils.InferenceResponse(
# output_tensors=..., TritonError("An error occurred"))
inference_response = pb_utils.InferenceResponse(
output_tensors=[output_tensor]
)
responses.append(inference_response)
# You should return a list of pb_utils.InferenceResponse. Length
# of this list must match the length of `requests` list.
return responses
def finalize(self) -> None:
"""Clean up the model.
`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 _postprocessing(self, tokens_batch: np.ndarray) -> List[bytes]:
"""Postprocess."""
outputs = []
for beam_tokens in tokens_batch:
for tokens in beam_tokens:
outputs.append(self.tokenizer.decode(tokens))
return outputs
|