Update app.py

app.py

@@ -52,9 +52,11 @@ def reformat_math(text):
     return text
 
 
-def user_input(message, history: list):
+def user_input(message, history_original, history_thinking):
     """사용자 입력을 히스토리에 추가하고 입력 텍스트 상자 비우기"""
-    return "", history + [
+    return "", history_original + [
+        gr.ChatMessage(role="user", content=message.replace(ANSWER_MARKER, ""))
+    ], history_thinking + [
         gr.ChatMessage(role="user", content=message.replace(ANSWER_MARKER, ""))
     ]
 
@@ -75,14 +77,63 @@ def rebuild_messages(history: list):
 
 
 @spaces.GPU
-def bot(
+def bot_original(
+    history: list,
+    max_num_tokens: int,
+    do_sample: bool,
+    temperature: float,
+):
+    """원본 모델이 질문에 답변하도록 하기 (추론 과정 없이)"""
+
+    # 나중에 스레드에서 토큰을 스트림으로 가져오기 위함
+    streamer = transformers.TextIteratorStreamer(
+        pipe.tokenizer,  # pyright: ignore
+        skip_special_tokens=True,
+        skip_prompt=True,
+    )
+
+    # 보조자 메시지 준비
+    history.append(
+        gr.ChatMessage(
+            role="assistant",
+            content=str(""),
+        )
+    )
+
+    # 현재 채팅에 표시될 메시지
+    messages = rebuild_messages(history[:-1])  # 마지막 빈 메시지 제외
+    
+    # 원본 모델은 추론 없이 바로 답변
+    t = threading.Thread(
+        target=pipe,
+        args=(messages,),
+        kwargs=dict(
+            max_new_tokens=max_num_tokens,
+            streamer=streamer,
+            do_sample=do_sample,
+            temperature=temperature,
+        ),
+    )
+    t.start()
+
+    for token in streamer:
+        history[-1].content += token
+        history[-1].content = reformat_math(history[-1].content)
+        yield history
+    t.join()
+
+    yield history
+
+
+@spaces.GPU
+def bot_thinking(
     history: list,
     max_num_tokens: int,
     final_num_tokens: int,
     do_sample: bool,
     temperature: float,
 ):
-    """모델이 질문에 답변하도록 하기"""
+    """추론 과정을 포함하여 모델이 질문에 답변하도록 하기"""
 
     # 나중에 스레드에서 토큰을 스트림으로 가져오기 위함
     streamer = transformers.TextIteratorStreamer(
@@ -142,21 +193,35 @@ def bot(
 
 with gr.Blocks(fill_height=True, title="모든 LLM 모델에 추론 능력 부여하기") as demo:
     with gr.Row(scale=1):
-        with gr.Column(scale=5):
-
-            chatbot = gr.Chatbot(
+        with gr.Column(scale=2):
+            gr.Markdown("## Before (Original)")
+            chatbot_original = gr.Chatbot(
                 scale=1,
                 type="messages",
                 latex_delimiters=latex_delimiters,
+                label="Original Model (No Reasoning)"
             )
-            msg = gr.Textbox(
-                submit_btn=True,
-                label="",
-                show_label=False,
-                placeholder="여기에 질문을 입력하세요.",
-                autofocus=True,
+        
+        with gr.Column(scale=2):
+            gr.Markdown("## After (Thinking)")
+            chatbot_thinking = gr.Chatbot(
+                scale=1,
+                type="messages",
+                latex_delimiters=latex_delimiters,
+                label="Model with Reasoning"
             )
-        with gr.Column(scale=1):
+    
+    with gr.Row():
+        msg = gr.Textbox(
+            submit_btn=True,
+            label="",
+            show_label=False,
+            placeholder="여기에 질문을 입력하세요.",
+            autofocus=True,
+        )
+    
+    with gr.Row():
+        with gr.Column():
             gr.Markdown("""## 매개변수 조정""")
             num_tokens = gr.Slider(
                 50,
@@ -177,22 +242,30 @@ with gr.Blocks(fill_height=True, title="모든 LLM 모델에 추론 능력 부
             do_sample = gr.Checkbox(True, label="샘플링 사용")
             temperature = gr.Slider(0.1, 1.0, 0.7, step=0.1, label="온도")
 
-
-    # 사용자가 메시지를 제출하면 봇이 응답합니다
+    # 사용자가 메시지를 제출하면 두 봇이 동시에 응답합니다
     msg.submit(
         user_input,
-        [msg, chatbot],  # 입력
-        [msg, chatbot],  # 출력
+        [msg, chatbot_original, chatbot_thinking],  # 입력
+        [msg, chatbot_original, chatbot_thinking],  # 출력
+    ).then(
+        bot_original,
+        [
+            chatbot_original,
+            num_tokens,
+            do_sample,
+            temperature,
+        ],
+        chatbot_original,  # 출력에서 새 히스토리 저장
     ).then(
-        bot,
+        bot_thinking,
         [
-            chatbot,
+            chatbot_thinking,
             num_tokens,
             final_num_tokens,
             do_sample,
             temperature,
-        ],  # 실제로는 "history" 입력
-        chatbot,  # 출력에서 새 히스토리 저장
+        ],
+        chatbot_thinking,  # 출력에서 새 히스토리 저장
     )
 
 if __name__ == "__main__":