add triton files

1/model.py

@@ -0,0 +1,234 @@
+import triton_python_backend_utils as pb_utils
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+import numpy as np
+import json
+import os
+
+class TritonPythonModel:
+    def initialize(self, args):
+        """
+        모델이 로드될 때 딱 한 번만 호출됩니다.
+        `initialize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
+        이 모델과 관련된 모든 상태를 초기화할 수 있습니다.
+        """
+        self.logger = pb_utils.Logger
+
+        current_file_path = os.path.abspath(__file__)
+        self.logger.log_info(f"current_file_path: {current_file_path}")
+
+
+        self.model_name = args["model_name"]
+        model_repository = args["model_repository"]
+        model_path = model_repository
+
+        self.logger.log_info(f"model_name: {self.model_name}")
+        self.logger.log_info(f"model_repository: {model_repository}")
+        self.logger.log_info(f"model_path: {model_path}")
+
+        self.model_config = json.loads(args["model_config"])
+
+        # Hugging Face Transformers 라이브러리에서 사전 학습된 토크나이저를 로드합니다.
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.supports_chat_template = self._check_chat_template_support()
+
+        # Hugging Face Transformers 라이브러리에서 사전 학습된 언어 모델을 로드합니다.
+        self.model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path=model_path,
+            local_files_only=True,
+            device_map="auto"
+        )
+
+        self.enable_inference_trace = self._get_inference_trace_setting()
+
+        self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료")
+
+
+    def execute(self, requests):
+        """
+        Triton이 각 추론 요청에 대해 호출하는 실행 함수입니다.
+        """
+        responses = []
+
+        # 각 추론 요청을 순회하며 처리합니다.
+        for request in requests:
+            # Triton 입력 파싱            
+            input_text = self._get_input_value(request, "text_input")
+            
+            text = ""
+            conversation = ""
+            input_token_length = 0  # 입력 토큰 길이를 저장할 변수
+
+            # 입력 텍스트가 JSON 형식의 대화 기록인지 확인합니다.
+            try:
+                conversation = json.loads(input_text)
+                is_chat = True
+                self.logger.log_info(f"입력 conversation 출력:\n{conversation}")
+            except:
+                # JSON 파싱에 실패하면 일반 텍스트로 처리합니다.
+                text = input_text
+                is_chat = False
+                self.logger.log_info(f"입력 text 출력:\n{text}")
+            
+            # 입력 텍스트를 토큰화합니다.
+            if self.supports_chat_template and is_chat:
+                self.logger.log_info(f"Chat 템플릿을 적용하여 토큰화합니다.")
+                inputs = self.tokenizer.apply_chat_template(
+                                            conversation,
+                                            tokenize=True,
+                                            add_generation_prompt=True,
+                                            return_tensors="pt",
+                                            return_dict=True
+                                            ).to(device=self.model.device)
+            else:
+                self.logger.log_info(f"입력 텍스트를 토큰화합니다.")
+                inputs = self.tokenizer(
+                            text,
+                            return_tensors="pt").to(device=self.model.device)
+
+            input_ids = inputs["input_ids"]
+            attention_mask = inputs["attention_mask"]
+            input_token_length = inputs["input_ids"].shape[-1]
+
+
+            # 언어 모델을 사용하여 텍스트를 생성합니다.
+            gened = self.model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                generation_config=self._process_generation_config(request),
+                pad_token_id=self.tokenizer.pad_token_id,
+            )
+
+            # 생성된 토큰 시퀀스를 텍스트로 디코딩하고 입력 텍스트는 제외합니다.
+            generated_tokens = gened[0][input_token_length:] # 입력 토큰 이후부터 슬라이싱
+            gened_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+            self.logger.log_info(f"모델이 생성한 토큰 시퀀스 (입력 텍스트 제외):\n{gened_text}")
+
+            output = gened_text.strip()
+
+            # 생성된 텍스트를 Triton 출력 텐서로 변환합니다.
+            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]))
+
+        return responses
+    
+    def _process_generation_config(self, request):
+        """
+        추론 요청에서 생성 설정 관련 파라미터들을 추출하여 GenerationConfig 객체를 생성합니다.
+
+        Args:
+            request (pb_utils.InferenceRequest): Triton 추론 요청 객체.
+
+        Returns:
+            transformers.GenerationConfig: GenerationConfig 객체.
+        """
+        max_length = self._get_input_value(request, "max_length", default=20)
+        max_new_tokens = self._get_input_value(request, "max_new_tokens")
+        temperature = self._get_input_value(request, "temperature")
+        do_sample = self._get_input_value(request, "do_sample")
+        top_k = self._get_input_value(request, "top_k")
+        top_p = self._get_input_value(request, "top_p")
+        repetition_penalty = self._get_input_value(request, "repetition_penalty")
+        stream = self._get_input_value(request, "stream")
+
+        generation_config = GenerationConfig(
+            max_length=max_length,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=do_sample,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            stream=stream,
+        )
+
+        self.logger.log_info(f"추론 요청 GenerationConfig:\n{generation_config}")
+
+        return generation_config
+
+    def _get_inference_trace_setting(self):
+        """
+        모델 설정(config.pbxt)에서 'enable_inference_trace' 값을 추출하여 반환합니다.
+
+        'enable_inference_trace' 설정이 없거나, 올바른 형식이 아닌 경우 기본적으로 False를 반환합니다.
+
+        Returns:
+            bool: 추론 추적 활성화 여부 (True 또는 False).
+        """
+        parameters = self.model_config.get('parameters', {})
+        trace_config = parameters.get('enable_inference_trace')
+        if isinstance(trace_config, dict) and 'string_value' in trace_config:
+            return trace_config['string_value'].lower() == 'true'  # 문자열 값을 bool로 변환하여 반환
+        return False
+    
+
+    def _check_chat_template_support(self):
+      """
+      주어진 허깅페이스 Transformer 모델이 Chat 템플릿을 지원하는지 확인하고 결과를 출력합니다.
+      
+      Returns:
+            bool: Chat 템플릿 지원 여부 (True 또는 False).
+      """
+      try:
+          if hasattr(self.tokenizer, "chat_template") and self.tokenizer.chat_template is not None:
+              self.logger.log_info(f"'{self.model_name}' 모델의 토크나이저는 Chat 템플릿을 지원합니다.")
+              self.logger.log_info("Chat 템플릿 내용:")
+              self.logger.log_info(self.tokenizer.chat_template)
+              return True
+          else:
+              self.logger.log_info(f"'{self.model_name}' 모델의 토크나이저는 Chat 템플릿을 직접적으로 지원하지 않거나, Chat 템플릿 정보가 없습니다.")
+              return False
+      except Exception as e:
+          self.logger.log_info(f"'{self.model_name}' 모델의 토크나이저를 로드하는 동안 오류가 발생했습니다: {e}")
+          return False
+    
+
+    def _get_input_value(self, request, input_name: str, default=None):
+        """
+        Triton 추론 요청에서 특정 이름의 입력 텐서 값을 가져옵니다.
+
+        Args:
+            request (pb_utils.InferenceRequest): Triton 추론 요청 객체.
+            input_name (str): 가져올 입력 텐서의 이름.
+            default (any, optional): 입력 텐서가 없을 경우 반환할 기본값. Defaults to None.
+
+        Returns:
+            any: 디코딩된 입력 텐서의 값. 텐서가 없으면 기본값을 반환합니다.
+        """
+        tensor_value = pb_utils.get_input_tensor_by_name(request, input_name)
+
+        if tensor_value is None:
+            return default
+        
+        return self._np_decoder(tensor_value.as_numpy()[0])
+        
+    def _np_decoder(self, obj):
+        """
+        NumPy 객체의 데이터 타입을 확인하고 Python 기본 타입으로 변환합니다.
+
+        Args:
+            obj (numpy.ndarray element): 변환할 NumPy 배열의 요소.
+
+        Returns:
+            any: 해당 NumPy 요소에 대응하는 Python 기본 타입 (str, int, float, bool).
+                 bytes 타입인 경우 UTF-8로 디코딩합니다.
+        """
+        if isinstance(obj, bytes):
+           return obj.decode('utf-8')
+        if np.issubdtype(obj, np.integer):
+           return int(obj)
+        if np.issubdtype(obj, np.floating):
+           return round(float(obj), 3)
+        if isinstance(obj, np.bool_):
+           return bool(obj)
+
+    def finalize(self):
+        """
+        모델 실행이 완료된 후 Triton 서버가 종료될 때 한 번 호출되는 함수입니다.
+        `finalize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
+        종료 전에 필요한 모든 정리 작업을 수행할 수 있습니다.
+        """
+        pass
+

config.pbtxt

@@ -0,0 +1,82 @@
+# Triton backend to use
+name: "base-gemma-3-1b-it"
+backend: "python"
+max_batch_size: 0
+
+# Triton should expect as input a single string
+# input of variable length named 'text_input'
+input [
+  {
+    name: "text_input"
+    data_type: TYPE_STRING
+    dims: [ -1 ]
+  },
+  {
+    name: "max_length"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "max_new_tokens"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "do_sample"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "top_k"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "top_p"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "temperature"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "repetition_penalty"
+    data_type: TYPE_FP32
+    dims: [ 1 ]
+    optional: true
+  },
+  {
+    name: "stream"
+    data_type: TYPE_BOOL
+    dims: [ 1 ]
+    optional: true
+  }
+]
+
+
+# Triton should expect to respond with a single string
+# output of variable length named 'text_output'
+output [ 
+  {
+    name: "text_output"
+    data_type: TYPE_STRING
+    dims: [ -1 ]
+  }
+]
+
+
+parameters: [
+  {
+    key: "enable_inference_trace",
+    value: {string_value: "False"}
+  }
+]
+