import torch import numpy as np import json 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): """ 모델이 로드될 때 딱 한 번만 호출됩니다. """ self.logger = pb_utils.Logger self.model_config = json.loads(args["model_config"]) 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 # 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 # 필수 옵션 추가 ) # 모델을 평가 모드로 설정 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 # 필수 옵션 추가 ) # 임베딩 모델이므로 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)") def execute(self, requests): """ Triton이 각 추론 요청에 대해 호출하는 실행 함수입니다. Generation 대신 Embedding 생성을 수행하도록 수정합니다. """ responses = [] # 각 추론 요청을 순회하며 처리합니다. for request in requests: # Triton 입력 파싱: 텍스트 입력만 처리합니다. input_text = self._get_input_value(request, "text_input") # CodeSage는 대화 형식이 아닌 일반 텍스트 (코드)를 입력으로 받으므로 # JSON 파싱 로직과 Chat 템플릿 로직을 제거합니다. text = input_text self.logger.log_info(f"입력 text 출력:\n{text}") # 입력 텍스트를 토큰화합니다. # add_eos_token=True가 load_model에서 설정되었으므로 토큰화 시 자동으로 추가됩니다. inputs = self.tokenizer( text, return_tensors="pt").to(device=self.model.device) input_ids = inputs["input_ids"] attention_mask = inputs["attention_mask"] # # CodeSage는 텍스트 생성이 아닌 임베딩 생성을 수행합니다. # gened = self.model.generate(...) # **임베딩 생성** # CodeSage 모델은 (임베딩, 히든 스테이트, 어텐션)을 반환하며, 첫 번째 요소가 임베딩입니다. with torch.no_grad(): # inputs에는 input_ids와 attention_mask가 모두 포함되어 전달됩니다. outputs = self.model(**inputs) # outputs[0]에는 임베딩 벡터가 포함되어 있습니다. # 임베딩은 일반적으로 첫 번째 토큰 (CLS 토큰 또는 문맥 임베딩)을 사용합니다. # CodeSage의 경우, 모델 카드 예시를 보면 outputs[0] 전체를 사용합니다. # 여기서는 [batch_size, sequence_length, hidden_size] 형태의 임베딩 중 첫 번째 토큰 임베딩을 사용합니다. # 임베딩 사용법은 모델의 목적에 따라 달라질 수 있습니다. CodeSage는 주로 전체 시퀀스 임베딩을 사용합니다. # 여기서는 예시와 같이 첫 번째 요소 (last_hidden_state)를 가져옵니다. # 임베딩 크기: [1, seq_len, hidden_size] embeddings = outputs[0] # 임베딩을 NumPy 배열로 변환 # CPU로 옮기고, NumPy로 변환 embedding_np = embeddings.squeeze().cpu().numpy() # 출력 텐서 생성 (데이터 타입은 float32 또는 float16이 적합) # CodeSage는 단일 문장 입력만 처리하므로, 배치 차원 없이 [seq_len, hidden_size]로 가정합니다. # 실제 사용 목적에 따라 풀링 로직을 추가하여 [hidden_size] 벡터로 만들 수도 있습니다. output_tensor = pb_utils.Tensor("embedding_output", embedding_np.astype(np.float32)) self.logger.log_info(f"모델이 생성한 임베딩 Shape:\n{embedding_np.shape}") # 응답 객체를 생성하고 출력 텐서를 추가합니다. 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): """ 모델 설정(config.pbtxt)에서 특정 파라미터의 문자열 값을 가져옵니다. """ self.parameters = self.model_config.get('parameters', {}) parameter_dict = self.parameters.get(parameter_name) if isinstance(parameter_dict, dict) and 'string_value' in parameter_dict: return parameter_dict['string_value'] return None def _get_input_value(self, request, input_name: str, default=None): """ Triton 추론 요청에서 특정 이름의 입력 텐서 값을 가져옵니다. """ 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 기본 타입으로 변환합니다. """ 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 서버가 종료될 때 한 번 호출되는 함수입니다. """ pass