Update 1/model.py

1/model.py

@@ -1,47 +1,41 @@
 import triton_python_backend_utils as pb_utils
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import numpy as np
-import random
-import string
 import json
 
 class TritonPythonModel:
-    """
-    Triton Python Model 클래스.
-    """
+    def initialize(self, args):
+        """
+        모델이 로드될 때 딱 한 번만 호출됩니다.
+        `initialize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
+        이 모델과 관련된 모든 상태를 초기화할 수 있습니다.
+        """
+        self.logger = pb_utils.Logger
+
+        self.model_name = args["model_name"]
+        self.model_config = json.loads(args["model_config"])
+
+        self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료")
 
-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")
-    
-    # Triton 데이터 타입 문자열을 NumPy 데이터 타입으로 직접 변환
-    # 'BYTES'는 np.object_ 타입에 해당함
-    if output_config['data_type'] == 'TYPE_STRING':
-        self.output_dtype = np.object_
-    else:
-        # 다른 데이터 타입에 대한 처리 로직 추가 가능
-        self.output_dtype = pb_utils.triton_string_to_np_dtype(output_config['data_type'])
 
     def execute(self, requests):
         """
-        클라이언트의 요청이 들어올 때마다 호출됩니다.
+        Triton이 각 추론 요청에 대해 호출하는 실행 함수입니다.
         """
         responses = []
 
+        # 각 추론 요청을 순회하며 처리합니다.
         for request in requests:
-            # 입력 데이터 가져오기 (이 예제에서는 입력 데이터를 사용하지 않음)
-            in_0 = pb_utils.get_input_tensor_by_name(request, "INPUT")
+            # Triton 입력 파싱            
+            input_text = self._get_input_value(request, "INPUT")
+            self.logger.log_info(f"INPUT 출력:\n{input_text}")
             
-            # 무작위 문자열 생성
-            random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))
-
-            output = random_string
+            
+            output = "HELLO_WORLD"
+            self.logger.log_info(f"OUTPUT 출력:\n{output}")
 
             # 생성된 텍스트를 Triton 출력 텐서로 변환합니다.
-            output_tensor = pb_utils.Tensor("OUTPUT", np.array(output.encode('utf-8'), dtype=np.bytes_))
+            output_tensor = pb_utils.Tensor("text_output", np.array(output.encode('utf-8'), dtype=np.bytes_))
 
             # 응답 객체를 생성하고 출력 텐서를 추가합니다.
             responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor]))
@@ -50,6 +44,8 @@ def initialize(self, args):
 
     def finalize(self):
         """
-        모델이 언로드될 때 호출됩니다.
+        모델 실행이 완료된 후 Triton 서버가 종료될 때 한 번 호출되는 함수입니다.
+        `finalize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
+        종료 전에 필요한 모든 정리 작업을 수행할 수 있습니다.
         """
-        print("TritonPythonModel: finalize() called.")
+        pass