diff --git a/1/model.py b/1/model.py index 82c8d5a..6a0787d 100644 --- a/1/model.py +++ b/1/model.py @@ -1,13 +1,11 @@ import torch -import numpy as np import json +import numpy as np import triton_python_backend_utils as pb_utils -# AutoModelForCausalLM 대신 AutoModel을 가져옵니다. + from transformers import AutoModel, AutoTokenizer, GenerationConfig from peft import PeftModel, PeftConfig -# BitsAndBytesConfig가 import 되지 않아 임시로 주석 처리하거나, 필요하다면 설치 후 주석 해제해야 합니다. -# from bitsandbytes import BitsAndBytesConfig class TritonPythonModel: def initialize(self, args): @@ -21,81 +19,30 @@ class TritonPythonModel: self.model_name = args["model_name"] self.base_model_path = self._get_config_parameter("base_model_path") - # CodeSage는 임베딩 모델이므로 LoRA 등의 어댑터 로드는 지원하지 않거나 일반적이지 않습니다. - # 기존 로직은 유지하되, 실제로 사용하지 않을 경우 config.pbtxt에서 해당 파라미터를 제거하는 것이 좋습니다. - self.is_adapter_model = self._get_config_parameter("is_adapter_model").strip().lower() == "true" - self.adapter_model_path = self._get_config_parameter("adapter_model_path") - self.quantization = self._get_config_parameter("quantization") - self.logger.log_info(f"base_model_path: {self.base_model_path}") - self.logger.log_info(f"is_adapter_model: {self.is_adapter_model}") - self.logger.log_info(f"adapter_model_path: {self.adapter_model_path}") - self.logger.log_info(f"quantization: {self.quantization}") + self.load_model() def load_model(self): - """ - Load model: CodeSage에 맞게 AutoModel과 trust_remote_code=True를 사용하도록 수정 - """ - self.bnb_config = None - torch_dtype = torch.float16 # 기본 dtype + torch_dtype = torch.float16 - # Note: CodeSage는 AutoModel을 사용하며, 일반적인 CausalLM 양자화 옵션이 적용되는지 - # 확인이 필요하지만, 일단 기존 로직은 유지합니다. - - # BitsAndBytesConfig가 정의되지 않았으므로 주석 처리합니다. - # if self.quantization == "int4": - # # ... (int4 config) - # pass - # elif self.quantization == "int8": - # # ... (int8 config) - # pass - # else: - # self.bnb_config = None - - if self.is_adapter_model: - # CodeSage는 임베딩 모델이므로 어댑터 사용은 일반적이지 않으나, - # 기존 템플릿 로직은 유지합니다. - peft_config = PeftConfig.from_pretrained(self.adapter_model_path) - self.base_model_path = peft_config.base_model_name_or_path - - # base 모델 로드: AutoModel로 변경 - base_model = AutoModel.from_pretrained( - peft_config.base_model_name_or_path, - torch_dtype=torch.float16, - quantization_config=self.bnb_config if self.bnb_config else None, - device_map="auto", - trust_remote_code=True # 필수 옵션 추가 - ) - - # adapter 모델 로드 - self.model = PeftModel.from_pretrained(base_model, self.adapter_model_path) - else: - # 일반 모델인 경우 로드: AutoModel로 변경 - self.model = AutoModel.from_pretrained( - pretrained_model_name_or_path=self.base_model_path, - # local_files_only=True, # CodeSage는 허브에서 로드될 수 있으므로 주석 처리 - quantization_config=self.bnb_config if self.bnb_config else None, - device_map="auto", - trust_remote_code=True # 필수 옵션 추가 - ) + # AutoModel로 모델 로드 + self.model = AutoModel.from_pretrained( + pretrained_model_name_or_path=self.base_model_path, + local_files_only=True, + device_map="auto", + trust_remote_code=False + ) # 모델을 평가 모드로 설정 self.model.eval() - # Tokenizer 로드: trust_remote_code=True와 add_eos_token=True 추가 self.tokenizer = AutoTokenizer.from_pretrained( self.base_model_path, - trust_remote_code=True, - add_eos_token=True # 필수 옵션 추가 + trust_remote_code=False, + add_eos_token=True ) - # 임베딩 모델이므로 pad_token_id 설정은 불필요하거나, 주의가 필요함. 일단 주석 처리 - # self.tokenizer.pad_token_id = self.tokenizer.eos_token_id - - # 임베딩 모델에서는 Chat Template 지원 체크는 불필요하므로 제거 - # self.supports_chat_template = self._check_chat_template_support() - self.supports_chat_template = False # 항상 False로 설정 self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료 (Code Embedding Mode)") @@ -158,23 +105,7 @@ class TritonPythonModel: responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor])) return responses - - # 임베딩 모델에서는 불필요하므로 제거합니다. - # def _process_generation_config(self, request): - # """ - # 추론 요청에서 생성 설정 관련 파라미터들을 추출하여 GenerationConfig 객체를 생성합니다. - # """ - # # ... (기존 로직) - # pass - # 임베딩 모델에서는 불필요하므로 제거합니다. - # def _check_chat_template_support(self): - # """ - # 주어진 허깅페이스 Transformer 모델이 Chat 템플릿을 지원하는지 확인하고 결과를 출력합니다. - # """ - # # ... (기존 로직) - # pass - def _get_config_parameter(self, parameter_name): """