Update 1/model.py
This commit is contained in:
parent
2ecddff2d3
commit
2e3e7382cc
83
1/model.py
83
1/model.py
@ -2,6 +2,7 @@ import json
|
||||
import torch
|
||||
import numpy as np
|
||||
import triton_python_backend_utils as pb_utils
|
||||
import uuid # [추가] 랜덤 ID 생성을 위해 필요
|
||||
from typing import List, Dict, Any, Union, Tuple
|
||||
|
||||
from transformers import (
|
||||
@ -44,7 +45,6 @@ class TritonPythonModel:
|
||||
bnb_config = self._get_bnb_config()
|
||||
|
||||
# 2. Base Model 로드
|
||||
# Adapter 모델인 경우 Config에서 Base 경로를 덮어쓸 수 있음
|
||||
load_path = self.base_model_path
|
||||
if self.is_adapter_model:
|
||||
peft_config = PeftConfig.from_pretrained(self.adapter_model_path)
|
||||
@ -67,26 +67,22 @@ class TritonPythonModel:
|
||||
if self.is_adapter_model:
|
||||
self.model = PeftModel.from_pretrained(self.model, self.adapter_model_path)
|
||||
|
||||
# 추론 모드 설정
|
||||
self.model.eval()
|
||||
|
||||
# 4. Tokenizer 로드
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(load_path, trust_remote_code=True)
|
||||
|
||||
# Pad Token 설정 (없을 경우 EOS로 대체)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
|
||||
self.logger.log_info("Pad token was None. Set to EOS token.")
|
||||
|
||||
# Chat Template 지원 여부 확인
|
||||
self.supports_chat_template = (
|
||||
hasattr(self.tokenizer, "chat_template") and
|
||||
self.tokenizer.chat_template is not None
|
||||
)
|
||||
|
||||
def _get_bnb_config(self) -> Union[BitsAndBytesConfig, None]:
|
||||
"""양자화 설정 객체를 반환합니다."""
|
||||
if self.quantization == "int4":
|
||||
return BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
@ -107,10 +103,19 @@ class TritonPythonModel:
|
||||
responses = []
|
||||
|
||||
for request in requests:
|
||||
# [ID 생성 로직]
|
||||
# 클라이언트가 보낸 ID가 있으면 사용하고, 없으면 랜덤 UUID 생성
|
||||
request_id = request.request_id()
|
||||
if not request_id:
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# 1. 입력 데이터 파싱
|
||||
input_data, is_chat = self._parse_input(request)
|
||||
self.logger.log_info(f"\n>>> [Input]: {input_data}")
|
||||
|
||||
# [LOGGING] Request ID 포함하여 로그 출력
|
||||
log_input_str = json.dumps(input_data, ensure_ascii=False) if isinstance(input_data, (list, dict)) else str(input_data)
|
||||
self.logger.log_info(f"\n[RID: {request_id}] >>> [Input]: {log_input_str}")
|
||||
|
||||
# 2. Generation Config 생성
|
||||
gen_config = self._create_generation_config(request)
|
||||
@ -121,37 +126,31 @@ class TritonPythonModel:
|
||||
# 4. 모델 추론 (Generate)
|
||||
output_text = self._generate(inputs, gen_config)
|
||||
|
||||
# [LOGGING] 모델 응답 로그 출력
|
||||
self.logger.log_info(f"\n<<< [Output]: {output_text}")
|
||||
# [LOGGING] Request ID 포함하여 결과 출력
|
||||
self.logger.log_info(f"\n[RID: {request_id}] <<< [Output]: {output_text}")
|
||||
|
||||
# 5. 응답 생성
|
||||
responses.append(self._create_response(output_text))
|
||||
# 5. 응답 생성 (ID 전달)
|
||||
responses.append(self._create_response(output_text, request_id))
|
||||
|
||||
except Exception as e:
|
||||
self.logger.log_error(f"Error during execution: {e}")
|
||||
# 에러 발생 시 빈 문자열 또는 에러 메시지 반환 (클라이언트 처리에 따라 변경 가능)
|
||||
self.logger.log_error(f"[RID: {request_id}] Error during execution: {e}")
|
||||
err_tensor = pb_utils.Tensor("text_output", np.array([str(e).encode('utf-8')], dtype=np.bytes_))
|
||||
responses.append(pb_utils.InferenceResponse(output_tensors=[err_tensor]))
|
||||
# 에러 응답에도 ID 포함
|
||||
responses.append(pb_utils.InferenceResponse(output_tensors=[err_tensor], id=request_id))
|
||||
|
||||
return responses
|
||||
|
||||
def _parse_input(self, request) -> Tuple[Union[str, List[Dict]], bool]:
|
||||
"""입력 텐서를 파싱하여 텍스트 또는 대화 목록과 타입(채팅 여부)을 반환합니다."""
|
||||
input_text = self._get_input_scalar(request, "text_input")
|
||||
|
||||
try:
|
||||
# JSON 형식의 대화 기록인지 시도
|
||||
conversation = json.loads(input_text)
|
||||
# 리스트 형태여야 채팅 히스토리로 간주
|
||||
if isinstance(conversation, list):
|
||||
return conversation, True
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
return input_text, False
|
||||
|
||||
def _tokenize(self, input_data, is_chat: bool):
|
||||
"""입력 데이터를 토큰화하여 모델 입력 텐서를 반환합니다."""
|
||||
if self.supports_chat_template and is_chat:
|
||||
return self.tokenizer.apply_chat_template(
|
||||
input_data,
|
||||
@ -161,17 +160,11 @@ class TritonPythonModel:
|
||||
return_dict=True
|
||||
).to(self.device)
|
||||
else:
|
||||
# 일반 텍스트인 경우
|
||||
if is_chat: # Chat 형식이지만 템플릿 미지원 시 문자열 변환 시도
|
||||
if is_chat:
|
||||
input_data = str(input_data)
|
||||
|
||||
return self.tokenizer(
|
||||
input_data,
|
||||
return_tensors="pt"
|
||||
).to(self.device)
|
||||
return self.tokenizer(input_data, return_tensors="pt").to(self.device)
|
||||
|
||||
def _generate(self, inputs, gen_config: GenerationConfig) -> str:
|
||||
"""모델 생성 로직 수행 및 디코딩"""
|
||||
input_ids = inputs["input_ids"]
|
||||
input_len = input_ids.shape[-1]
|
||||
|
||||
@ -183,15 +176,11 @@ class TritonPythonModel:
|
||||
eos_token_id=self.tokenizer.eos_token_id
|
||||
)
|
||||
|
||||
# 입력 토큰을 제외하고 생성된 토큰만 슬라이싱
|
||||
generated_tokens = outputs[0][input_len:]
|
||||
decoded_output = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
|
||||
|
||||
return decoded_output.strip()
|
||||
|
||||
def _create_generation_config(self, request) -> GenerationConfig:
|
||||
"""Request에서 파라미터를 추출하여 GenerationConfig 객체 생성"""
|
||||
# 기본값 설정 및 입력값 추출 Helper
|
||||
def get_param(name, default=None, cast_type=None):
|
||||
val = self._get_input_scalar(request, name, default)
|
||||
if val is not None and cast_type:
|
||||
@ -199,56 +188,37 @@ class TritonPythonModel:
|
||||
return val
|
||||
|
||||
return GenerationConfig(
|
||||
max_length=get_param("max_length", 1024, int), # max_length 기본값은 넉넉하게
|
||||
max_length=get_param("max_length", 1024, int),
|
||||
max_new_tokens=get_param("max_new_tokens", 256, int),
|
||||
temperature=get_param("temperature", 1.0, float),
|
||||
do_sample=get_param("do_sample", False, bool),
|
||||
top_k=get_param("top_k", 50, int),
|
||||
top_p=get_param("top_p", 1.0, float),
|
||||
repetition_penalty=get_param("repetition_penalty", 1.0, float),
|
||||
# stream=get_param("stream", False, bool) # Python Backend에서 Stream은 별도 구현 필요
|
||||
)
|
||||
|
||||
def _create_response(self, output_text: str):
|
||||
"""생성된 텍스트를 Triton Response 객체로 변환"""
|
||||
def _create_response(self, output_text: str, request_id: str):
|
||||
"""생성된 텍스트를 Triton Response 객체로 변환 (ID 포함)"""
|
||||
output_tensor = pb_utils.Tensor(
|
||||
"text_output",
|
||||
np.array([output_text.encode('utf-8')], dtype=np.bytes_)
|
||||
)
|
||||
return pb_utils.InferenceResponse(output_tensors=[output_tensor])
|
||||
# id 파라미터를 추가하여 클라이언트에게도 ID를 반환
|
||||
return pb_utils.InferenceResponse(output_tensors=[output_tensor], id=request_id)
|
||||
|
||||
def _get_config_param(self, key: str, default: str = None) -> str:
|
||||
"""config.pbtxt 파라미터 조회 Helper"""
|
||||
params = self.model_config.get('parameters', {})
|
||||
if key in params:
|
||||
return params[key].get('string_value', default)
|
||||
return default
|
||||
|
||||
def _get_input_scalar(self, request, name: str, default=None):
|
||||
"""입력 텐서에서 스칼라 값을 추출하는 Helper"""
|
||||
tensor = pb_utils.get_input_tensor_by_name(request, name)
|
||||
if tensor is None:
|
||||
return default
|
||||
|
||||
#value = tensor.as_numpy().item() # item()을 사용하여 스칼라 값 추출
|
||||
value = self._np_decoder(tensor.as_numpy()[0])
|
||||
|
||||
# 바이트 타입 디코딩
|
||||
# if isinstance(value, bytes):
|
||||
# return value.decode('utf-8')
|
||||
return value
|
||||
return self._np_decoder(tensor.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):
|
||||
@ -259,7 +229,6 @@ class TritonPythonModel:
|
||||
return bool(obj)
|
||||
|
||||
def finalize(self):
|
||||
"""리소스 정리"""
|
||||
self.logger.log_info(f"Finalizing model {self.model_name}")
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
|
||||
Loading…
Reference in New Issue
Block a user