import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer import gc import os import datetime import time import spaces # --- 설정 --- MODEL_ID = "naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B" MAX_NEW_TOKENS = 512 CPU_THREAD_COUNT = 4 # 필요시 조절 HF_TOKEN = os.getenv("HF_TOKEN") # Hugging Face 토큰 환경 변수에서 가져오기 # --- 선택 사항: CPU 스레드 설정 --- # torch.set_num_threads(CPU_THREAD_COUNT) # os.environ["OMP_NUM_THREADS"] = str(CPU_THREAD_COUNT) # os.environ["MKL_NUM_THREADS"] = str(CPU_THREAD_COUNT) print("--- 환경 설정 ---") print(f"PyTorch 버전: {torch.__version__}") print(f"실행 장치: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}") print(f"Torch 스레드: {torch.get_num_threads()}") print(f"HF_TOKEN 설정 여부: {'있음' if HF_TOKEN else '없음'}") # --- 모델 및 토크나이저 로딩 --- print(f"--- 모델 로딩 중: {MODEL_ID} ---") print("첫 실행 시 몇 분 정도 소요될 수 있습니다...") model = None tokenizer = None load_successful = False stop_token_ids_list = [] # stop_token_ids_list 초기화 try: start_load_time = time.time() # 자원에 따라 device_map 설정 device_map = "auto" if torch.cuda.is_available() else "cpu" dtype = torch.float16 if torch.cuda.is_available() else torch.float32 tokenizer = AutoTokenizer.from_pretrained( MODEL_ID, trust_remote_code=True, token=HF_TOKEN # 토큰 추가 ) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=dtype, device_map=device_map, trust_remote_code=True, token=HF_TOKEN # 토큰 추가 ) model.eval() load_time = time.time() - start_load_time print(f"--- 모델 및 토크나이저 로딩 완료: {load_time:.2f}초 소요 ---") load_successful = True # --- 중지 토큰 설정 --- stop_token_strings = ["", "<|endoftext|>"] temp_stop_ids = [tokenizer.convert_tokens_to_ids(token) for token in stop_token_strings] if tokenizer.eos_token_id is not None and tokenizer.eos_token_id not in temp_stop_ids: temp_stop_ids.append(tokenizer.eos_token_id) elif tokenizer.eos_token_id is None: print("경고: tokenizer.eos_token_id가 None입니다. 중지 토큰에 추가할 수 없습니다.") stop_token_ids_list = [tid for tid in temp_stop_ids if tid is not None] if not stop_token_ids_list: print("경고: 중지 토큰 ID를 찾을 수 없습니다. 가능하면 기본 EOS를 사용하고, 그렇지 않으면 생성이 올바르게 중지되지 않을 수 있습니다.") if tokenizer.eos_token_id is not None: stop_token_ids_list = [tokenizer.eos_token_id] else: print("오류: 기본 EOS를 포함하여 중지 토큰을 찾을 수 없습니다. 생성이 무한정 실행될 수 있습니다.") print(f"사용할 중지 토큰 ID: {stop_token_ids_list}") except Exception as e: print(f"!!! 모델 로딩 오류: {e}") if 'model' in locals() and model is not None: del model if 'tokenizer' in locals() and tokenizer is not None: del tokenizer gc.collect() raise gr.Error(f"모델 {MODEL_ID} 로딩에 실패했습니다. 애플리케이션을 시작할 수 없습니다. 오류: {e}") # --- 시스템 프롬프트 정의 --- def get_system_prompt(): current_date = datetime.datetime.now().strftime("%Y-%m-%d (%A)") return ( f"- 오늘은 {current_date}입니다.\n" f"- 사용자의 질문에 대해 친절하고 자세하게 한국어로 답변해야 합니다." ) # --- 웜업 함수 --- def warmup_model(): if not load_successful or model is None or tokenizer is None: print("웜업 건너뛰기: 모델이 성공적으로 로드되지 않았습니다.") return print("--- 모델 웜업 시작 ---") try: start_warmup_time = time.time() warmup_message = "안녕하세요" # 모델에 맞는 형식으로 입력 구성 system_prompt = get_system_prompt() # MiMo 모델의 프롬프트 형식에 맞게 조정 prompt = f"Human: {warmup_message}\nAssistant:" inputs = tokenizer(prompt, return_tensors="pt").to(model.device) # 중지 토큰이 비어 있는지 확인하고 적절히 처리 gen_kwargs = { "max_new_tokens": 10, "pad_token_id": tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id, "do_sample": False } if stop_token_ids_list: gen_kwargs["eos_token_id"] = stop_token_ids_list else: print("웜업 경고: 생성에 정의된 중지 토큰이 없습니다.") with torch.no_grad(): output_ids = model.generate(**inputs, **gen_kwargs) del inputs del output_ids gc.collect() warmup_time = time.time() - start_warmup_time print(f"--- 모델 웜업 완료: {warmup_time:.2f}초 소요 ---") except Exception as e: print(f"!!! 모델 웜업 중 오류 발생: {e}") finally: gc.collect() # --- 추론 함수 --- @spaces.GPU() def predict(message, history): """ HyperCLOVAX-SEED-Vision-Instruct-3B 모델을 사용하여 응답을 생성합니다. 'history'는 Gradio 'messages' 형식을 가정합니다: List[Dict]. """ if model is None or tokenizer is None: return "오류: 모델이 로드되지 않았습니다." # 대화 기록 처리 history_text = "" if isinstance(history, list): for turn in history: if isinstance(turn, tuple) and len(turn) == 2: history_text += f"Human: {turn[0]}\nAssistant: {turn[1]}\n" # MiMo 모델 입력 형식에 맞게 프롬프트 구성 prompt = f"{history_text}Human: {message}\nAssistant:" inputs = None output_ids = None try: # 입력 준비 inputs = tokenizer(prompt, return_tensors="pt").to(model.device) input_length = inputs.input_ids.shape[1] print(f"\n입력 토큰 수: {input_length}") except Exception as e: print(f"!!! 입력 처리 중 오류 발생: {e}") return f"오류: 입력 형식을 처리하는 중 문제가 발생했습니다. ({e})" try: print("응답 생성 중...") generation_start_time = time.time() # 생성 인수 준비, 비어 있는 stop_token_ids_list 처리 gen_kwargs = { "max_new_tokens": MAX_NEW_TOKENS, "pad_token_id": tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.pad_token_id, "do_sample": True, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.1 } if stop_token_ids_list: gen_kwargs["eos_token_id"] = stop_token_ids_list else: print("생성 경고: 정의된 중지 토큰이 없습니다.") with torch.no_grad(): output_ids = model.generate(**inputs, **gen_kwargs) generation_time = time.time() - generation_start_time print(f"생성 완료: {generation_time:.2f}초 소요.") except Exception as e: print(f"!!! 모델 생성 중 오류 발생: {e}") if inputs is not None: del inputs if output_ids is not None: del output_ids gc.collect() return f"오류: 응답을 생성하는 중 문제가 발생했습니다. ({e})" # 응답 디코딩 response = "오류: 응답 생성에 실패했습니다." if output_ids is not None: try: new_tokens = output_ids[0, input_length:] response = tokenizer.decode(new_tokens, skip_special_tokens=True) print(f"출력 토큰 수: {len(new_tokens)}") del new_tokens except Exception as e: print(f"!!! 응답 디코딩 중 오류 발생: {e}") response = "오류: 응답을 디코딩하는 중 문제가 발생했습니다." # 메모리 정리 if inputs is not None: del inputs if output_ids is not None: del output_ids gc.collect() print("메모리 정리 완료.") return response.strip() # --- Gradio 인터페이스 설정 --- print("--- Gradio 인터페이스 설정 중 ---") examples = [ ["안녕하세요! 자기소개 좀 해주세요."], ["인공지능과 머신러닝의 차이점은 무엇인가요?"], ["딥러닝 모델 학습 과정을 단계별로 알려주세요."], ["제주도 여행 계획을 세우고 있는데, 3박 4일 추천 코스 좀 알려주세요."], ] # ChatInterface를 사용하여 자체 Chatbot 컴포넌트 관리 demo = gr.ChatInterface( fn=predict, title="🤖 HyperCLOVAX-SEED-Text-Instruct-0.5B", description=( f"**모델:** {MODEL_ID}\n" ), examples=examples, cache_examples=False, theme=gr.themes.Soft(), ) # --- 애플리케이션 실행 --- if __name__ == "__main__": if load_successful: warmup_model() else: print("모델 로딩에 실패하여 웜업을 건너뜁니다.") print("--- Gradio 앱 실행 중 ---") demo.queue().launch( # share=True # 공개 링크를 원하면 주석 해제 # server_name="0.0.0.0" # 로컬 네트워크 접근을 원하면 주석 해제 )