import re import threading import gradio as gr import spaces import transformers from transformers import pipeline # 사용 가능한 모델 목록 available_models = { "meta-llama/Llama-3.2-3B-Instruct": "Llama 3.2(3B)", "Hermes-3-Llama-3.1-8B": "Hermes 3 Llama 3.1 (8B)", "nvidia/Llama-3.1-Nemotron-Nano-8B-v1": "Nvidia Nemotron Nano (8B)", "mistralai/Mistral-Small-3.1-24B-Instruct-2503": "Mistral Small 3.1 (24B)", "bartowski/mistralai_Mistral-Small-3.1-24B-Instruct-2503-GGUF": "Mistral Small GGUF (24B)", "google/gemma-3-27b-it": "Google Gemma 3 (27B)", "gemma-3-27b-it-abliterated": "Gemma 3 Abliterated (27B)", "Qwen/Qwen2.5-Coder-32B-Instruct": "Qwen 2.5 Coder (32B)", "open-r1/OlympicCoder-32B": "Olympic Coder (32B)" } # 모델과 토크나이저 로딩을 위한 전역 변수 pipe = None # 최종 답변을 감지하기 위한 마커 ANSWER_MARKER = "**답변**" # 단계별 추론을 시작하는 문장들 rethink_prepends = [ "자, 이제 다음을 파악해야 합니다 ", "제 생각에는 ", "잠시만요, 제 생각에는 ", "다음 사항이 맞는지 확인해 보겠습니다 ", "또한 기억해야 할 것은 ", "또 다른 주목할 점은 ", "그리고 저는 다음과 같은 사실도 기억합니다 ", "이제 충분히 이해했다고 생각합니다 ", "지금까지의 정보를 바탕으로, 원래 질문에 사용된 언어로 답변하겠습니다:" "\n{question}\n" f"\n{ANSWER_MARKER}\n", ] # 수식 표시 문제 해결을 위한 설정 latex_delimiters = [ {"left": "$$", "right": "$$", "display": True}, {"left": "$", "right": "$", "display": False}, ] def reformat_math(text): """Gradio 구문(Katex)을 사용하도록 MathJax 구분 기호 수정. 이것은 Gradio에서 수학 공식을 표시하기 위한 임시 해결책입니다. 현재로서는 다른 latex_delimiters를 사용하여 예상대로 작동하게 하는 방법을 찾지 못했습니다... """ text = re.sub(r"\\\[\s*(.*?)\s*\\\]", r"$$\1$$", text, flags=re.DOTALL) text = re.sub(r"\\\(\s*(.*?)\s*\\\)", r"$\1$", text, flags=re.DOTALL) return text def user_input(message, history: list): """사용자 입력을 히스토리에 추가하고 입력 텍스트 상자 비우기""" return "", history + [ gr.ChatMessage(role="user", content=message.replace(ANSWER_MARKER, "")) ] def rebuild_messages(history: list): """중간 생각 과정 없이 모델이 사용할 히스토리에서 메시지 재구성""" messages = [] for h in history: if isinstance(h, dict) and not h.get("metadata", {}).get("title", False): messages.append(h) elif ( isinstance(h, gr.ChatMessage) and h.metadata.get("title") and isinstance(h.content, str) ): messages.append({"role": h.role, "content": h.content}) return messages def load_model(model_names): """선택된 모델 이름에 따라 모델 로드""" global pipe # 모델이 선택되지 않았을 경우 기본값 지정 if not model_names: model_name = "Qwen/Qwen2-1.5B-Instruct" else: # 첫 번째 선택된 모델 사용 (나중에 여러 모델 앙상블로 확장 가능) model_name = model_names[0] pipe = pipeline( "text-generation", model=model_name, device_map="auto", torch_dtype="auto", ) return f"모델 '{model_name}'이(가) 로드되었습니다." @spaces.GPU def bot( history: list, max_num_tokens: int, final_num_tokens: int, do_sample: bool, temperature: float, ): """모델이 질문에 답변하도록 하기""" global pipe # 모델이 로드되지 않았다면 오류 메시지 표시 if pipe is None: history.append( gr.ChatMessage( role="assistant", content="모델이 로드되지 않았습니다. 하나 이상의 모델을 선택해 주세요.", ) ) yield history return # 나중에 스레드에서 토큰을 스트림으로 가져오기 위함 streamer = transformers.TextIteratorStreamer( pipe.tokenizer, # pyright: ignore skip_special_tokens=True, skip_prompt=True, ) # 필요한 경우 추론에 질문을 다시 삽입하기 위함 question = history[-1]["content"] # 보조자 메시지 준비 history.append( gr.ChatMessage( role="assistant", content=str(""), metadata={"title": "🧠 생각 중...", "status": "pending"}, ) ) # 현재 채팅에 표시될 추론 과정 messages = rebuild_messages(history) for i, prepend in enumerate(rethink_prepends): if i > 0: messages[-1]["content"] += "\n\n" messages[-1]["content"] += prepend.format(question=question) num_tokens = int( max_num_tokens if ANSWER_MARKER not in prepend else final_num_tokens ) t = threading.Thread( target=pipe, args=(messages,), kwargs=dict( max_new_tokens=num_tokens, streamer=streamer, do_sample=do_sample, temperature=temperature, ), ) t.start() # 새 내용으로 히스토리 재구성 history[-1].content += prepend.format(question=question) if ANSWER_MARKER in prepend: history[-1].metadata = {"title": "💭 사고 과정", "status": "done"} # 생각 종료, 이제 답변입니다 (중간 단계에 대한 메타데이터 없음) history.append(gr.ChatMessage(role="assistant", content="")) for token in streamer: history[-1].content += token history[-1].content = reformat_math(history[-1].content) yield history t.join() yield history with gr.Blocks(fill_height=True, title="ThinkFlow - Step-by-step Reasoning Service") as demo: # 상단에 타이틀과 설명 추가 gr.Markdown(""" # ThinkFlow ## A thought amplification service that implants step-by-step reasoning abilities into LLMs without model modification """) with gr.Row(scale=1): with gr.Column(scale=5): # 채팅 인터페이스 chatbot = gr.Chatbot( scale=1, type="messages", latex_delimiters=latex_delimiters, ) msg = gr.Textbox( submit_btn=True, label="", show_label=False, placeholder="여기에 질문을 입력하세요.", autofocus=True, ) with gr.Column(scale=1): # 모델 선택 섹션 추가 gr.Markdown("""## 모델 선택""") model_selector = gr.CheckboxGroup( choices=list(available_models.values()), value=[available_models["Qwen/Qwen2-1.5B-Instruct"]], # 기본값 label="사용할 LLM 모델 선택 (복수 선택 가능)", ) # 모델 로드 버튼 load_model_btn = gr.Button("모델 로드") model_status = gr.Textbox(label="모델 상태", interactive=False) gr.Markdown("""## 매개변수 조정""") num_tokens = gr.Slider( 50, 4000, 2000, step=1, label="추론 단계당 최대 토큰 수", interactive=True, ) final_num_tokens = gr.Slider( 50, 4000, 2000, step=1, label="최종 답변의 최대 토큰 수", interactive=True, ) do_sample = gr.Checkbox(True, label="샘플링 사용") temperature = gr.Slider(0.1, 1.0, 0.7, step=0.1, label="온도") # 선택된 모델 로드 이벤트 연결 def get_model_names(selected_models): # 표시 이름에서 원래 모델 이름으로 변환 inverse_map = {v: k for k, v in available_models.items()} return [inverse_map[model] for model in selected_models] load_model_btn.click( lambda selected: load_model(get_model_names(selected)), inputs=[model_selector], outputs=[model_status] ) # 사용자가 메시지를 제출하면 봇이 응답합니다 msg.submit( user_input, [msg, chatbot], # 입력 [msg, chatbot], # 출력 ).then( bot, [ chatbot, num_tokens, final_num_tokens, do_sample, temperature, ], # 실제로는 "history" 입력 chatbot, # 출력에서 새 히스토리 저장 ) if __name__ == "__main__": demo.queue().launch()