app.py

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()