Update app.py

app.py

@@ -6,104 +6,69 @@ from threading import Event
 hf_token = os.getenv("HF_TOKEN")
 stop_event = Event()
 
-def get_model_response(client, messages, max_tokens, temperature, top_p):
+models = {
+    "deepseek-ai/DeepSeek-Coder-V2-Instruct": "(한국회사)DeepSeek-Coder-V2-Instruct",
+    "meta-llama/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct",
+    "mistralai/Mixtral-8x7B-Instruct-v0.1": "Mixtral-8x7B-Instruct-v0.1",
+    "CohereForAI/c4ai-command-r-plus": "Cohere Command-R Plus"
+}
+
+def get_client(model):
+    return InferenceClient(model=model, token=hf_token)
+
+def respond(message, chat_history, system_message, max_tokens, temperature, top_p, selected_model):
+    stop_event.clear()
+    client = get_client(selected_model)
+    
+    messages = [{"role": "system", "content": system_message}]
+    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in chat_history for i, m in enumerate(h) if m])
+    messages.append({"role": "user", "content": message})
+    
     try:
-        response = client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
+        response = ""
+        for chunk in client.text_generation(
+            prompt="\n".join([f"{m['role']}: {m['content']}" for m in messages]),
+            max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             stream=True
-        )
-        full_response = ""
-        for message in response:
+        ):
             if stop_event.is_set():
                 break
-            if hasattr(message.choices[0], 'delta'):
-                token = message.choices[0].delta.content
-            else:
-                token = message.choices[0].text
-            if token:
-                full_response += token
-                yield full_response
-    except Exception as e:
-        yield f"모델 추론 실패: {str(e)}"
-
-def respond(message, history, system_message, max_tokens, temperature, top_p, selected_model):
-    stop_event.clear()
-    try:
-        client = InferenceClient(model=selected_model, token=hf_token)
-        
-        messages = [{"role": "system", "content": system_message}]
-        messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in history for i, m in enumerate(h) if m])
-        messages.append({"role": "user", "content": message})
-        
-        history.append((message, ""))
-        for response in get_model_response(client, messages, max_tokens, temperature, top_p):
-            history[-1] = (message, response)
-            yield "", history
-        
+            if chunk:
+                response += chunk
+                chat_history.append((message, response))
+                yield chat_history
     except Exception as e:
-        history.append((message, f"오류 발생: {str(e)}"))
-        yield "", history
+        error_message = f"오류 발생: {str(e)}"
+        chat_history.append((message, error_message))
+        yield chat_history
 
-def continue_writing(history, system_message, max_tokens, temperature, top_p, model):
-    if not history:
-        return "", history
-    last_user_message = history[-1][0]
-    last_assistant_message = history[-1][1]
-    
+def continue_writing(message, chat_history, system_message, max_tokens, temperature, top_p, selected_model):
+    if not chat_history:
+        return [("시스템", "대화 내역이 없습니다.")]
+    last_user_message = chat_history[-1][0]
+    last_assistant_message = chat_history[-1][1]
     prompt = f"이전 대화를 계속 이어서 작성해주세요. 이전 응답: {last_assistant_message}"
-    
-    try:
-        client = InferenceClient(model=model, token=hf_token)
-        messages = [{"role": "system", "content": system_message}]
-        messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for h in history for i, m in enumerate(h)])
-        messages.append({"role": "user", "content": prompt})
-        
-        for response in get_model_response(client, messages, max_tokens, temperature, top_p):
-            continued_response = last_assistant_message + " " + response
-            history[-1] = (last_user_message, continued_response)
-            yield "", history
-    except Exception as e:
-        history.append(("시스템", f"계속 작성 중 오류 발생: {str(e)}"))
-        yield "", history
+    return respond(prompt, chat_history, system_message, max_tokens, temperature, top_p, selected_model)
 
 def stop_generation():
     stop_event.set()
     return "생성이 중단되었습니다."
 
-def regenerate(history, system_message, max_tokens, temperature, top_p, model):
-    if not history:
-        return "", []
-    last_user_message = history[-1][0]
-    new_history = history[:-1]
-    for _, updated_history in respond(last_user_message, new_history, system_message, max_tokens, temperature, top_p, model):
-        yield "", updated_history
-
-models = {
-    "deepseek-ai/DeepSeek-Coder-V2-Instruct": "(한국회사)DeepSeek-Coder-V2-Instruct",
-    "meta-llama/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct",
-    "mistralai/Mixtral-8x7B-Instruct-v0.1": "Mixtral-8x7B-Instruct-v0.1",
-    "CohereForAI/c4ai-command-r-plus": "Cohere Command-R Plus"
-}
+def regenerate(chat_history, system_message, max_tokens, temperature, top_p, selected_model):
+    if not chat_history:
+        return [("시스템", "대화 내역이 없습니다.")]
+    last_user_message = chat_history[-1][0]
+    return respond(last_user_message, chat_history[:-1], system_message, max_tokens, temperature, top_p, selected_model)
 
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(height=600)
-    
-    with gr.Row():
-        msg = gr.Textbox(
-            scale=4,
-            label="메시지 입력",
-            lines=3,
-            placeholder="메시지를 입력하세요. Enter로 줄바꿈, Shift+Enter로 전송"
-        )
-    
-    with gr.Row():
-        send = gr.Button("전송", scale=1)
-        continue_btn = gr.Button("계속 작성", scale=1)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="메시지 입력", placeholder="메시지를 입력하세요. Enter로 전송, Shift+Enter로 줄바꿈")
 
     with gr.Row():
+        send = gr.Button("전송")
+        continue_btn = gr.Button("계속 작성")
         regenerate_btn = gr.Button("🔄 재생성")
         stop = gr.Button("🛑 생성 중단")
         clear = gr.Button("🗑️ 대화 내역 지우기")
@@ -112,35 +77,20 @@ with gr.Blocks() as demo:
         system_message = gr.Textbox(
             value="너는 나의 최고의 비서이다.\n내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.\n반드시 한글로 답변할것.",
             label="시스템 메시지",
-            lines=10
+            lines=5
         )
         max_tokens = gr.Slider(minimum=1, maximum=2000, value=500, step=100, label="최대 새 토큰 수")
         temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="온도")
         top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.90, step=0.05, label="Top-p (핵 샘플링)")
         model = gr.Radio(list(models.keys()), value=list(models.keys())[0], label="언어 모델 선택", info="사용할 언어 모델을 선택하세요")
 
-    msg.javascript = """
-    function(x) {
-        const textbox = document.getElementById('component-3');
-        if (!textbox) return x;
-        
-        textbox.addEventListener("keydown", function(e) {
-            if (e.key === 'Enter' && e.shiftKey) {
-                e.preventDefault();
-                const sendButton = document.getElementById('component-5');
-                if (sendButton) sendButton.click();
-            }
-        });
-        return x;
-    }
-    """
-
-    send.click(respond, inputs=[msg, chatbot, system_message, max_tokens, temperature, top_p, model], outputs=[msg, chatbot])
-    msg.submit(respond, inputs=[msg, chatbot, system_message, max_tokens, temperature, top_p, model], outputs=[msg, chatbot])
-    continue_btn.click(continue_writing, inputs=[chatbot, system_message, max_tokens, temperature, top_p, model], outputs=[msg, chatbot])
-    regenerate_btn.click(regenerate, inputs=[chatbot, system_message, max_tokens, temperature, top_p, model], outputs=[msg, chatbot])
-    stop.click(stop_generation, inputs=[], outputs=[msg])
-    clear.click(lambda: (None, None), outputs=[msg, chatbot])
+    # Event handlers
+    msg.submit(respond, [msg, chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
+    send.click(respond, [msg, chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
+    continue_btn.click(continue_writing, [msg, chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
+    regenerate_btn.click(regenerate, [chatbot, system_message, max_tokens, temperature, top_p, model], [chatbot])
+    stop.click(stop_generation, outputs=[msg])
+    clear.click(lambda: None, outputs=[chatbot])
 
 if __name__ == "__main__":
     if not hf_token: