triton-python-backend-model/1/model.py

54 lines
1.7 KiB
Python

import triton_python_backend_utils as pb_utils
import numpy as np
import random
import string
import json
class TritonPythonModel:
"""
Triton Python Model 클래스.
"""
def initialize(self, args):
"""
모델이 로드될 때 호출됩니다.
"""
print("TritonPythonModel: initialize() called.")
self.model_config = json.loads(args['model_config'])
output_config = pb_utils.get_output_config_by_name(
self.model_config, "OUTPUT")
self.output_dtype = pb_utils.triton_string_to_np_dtype(
output_config['data_type'])
def execute(self, requests):
"""
클라이언트의 요청이 들어올 때마다 호출됩니다.
"""
responses = []
for request in requests:
# 입력 데이터 가져오기 (이 예제에서는 입력 데이터를 사용하지 않음)
in_0 = pb_utils.get_input_tensor_by_name(request, "INPUT")
# 무작위 문자열 생성
random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))
# numpy 배열로 변환
output_array = np.array([[random_string.encode('utf-8')]], dtype=self.output_dtype)
# 출력 텐서 생성
out_tensor = pb_utils.Tensor("OUTPUT", output_array)
# 추론 응답 생성
inference_response = pb_utils.InferenceResponse(
output_tensors=[out_tensor])
responses.append(inference_response)
return responses
def finalize(self):
"""
모델이 언로드될 때 호출됩니다.
"""
print("TritonPythonModel: finalize() called.")