Update app.py

app.py

@@ -16,14 +16,9 @@ models = {
 def get_client(model):
     return InferenceClient(model=model, token=hf_token)
 
-MAX_HISTORY_LENGTH = 5  # 히스토리에 유지할 최대 대화 수
-
-def truncate_history(history):
-    return history[-MAX_HISTORY_LENGTH:] if len(history) > MAX_HISTORY_LENGTH else history
-
 def respond(message, system_message, max_tokens, temperature, top_p, selected_model):
     stop_event.clear()
-    client = InferenceClient(model=selected_model, token=hf_token)
+    client = get_client(selected_model)
     
     messages = [
         {"role": "system", "content": system_message},
@@ -48,23 +43,15 @@ def respond(message, system_message, max_tokens, temperature, top_p, selected_mo
     except Exception as e:
         yield [(message, f"오류 발생: {str(e)}")]
 
-def stop_generation():
-    stop_event.set()
-    return "생성이 중단되었습니다."
-    
-def stop_generation():
-    stop_event.set()
-    return "생성이 중단되었습니다."
-
-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)
+def get_last_response(chatbot):
+    if chatbot and len(chatbot) > 0:
+        return chatbot[-1][1]
+    return ""
 
-def continue_writing(last_response, system_message, max_tokens, temperature, top_p, selected_model):
+def continue_writing(chatbot, system_message, max_tokens, temperature, top_p, selected_model):
+    last_response = get_last_response(chatbot)
     stop_event.clear()
-    client = InferenceClient(model=selected_model, token=hf_token)
+    client = get_client(selected_model)
     
     prompt = f"이전 응답을 이어서 작성해주세요. 이전 응답: {last_response}"
     messages = [
@@ -85,12 +72,15 @@ def continue_writing(last_response, system_message, max_tokens, temperature, top
                 break
             if chunk:
                 response += chunk
-                yield [("계속 작성", response)]
+                yield chatbot + [("계속 작성", response)]
         
     except Exception as e:
-        yield [("계속 작성", f"오류 발생: {str(e)}")]
+        yield chatbot + [("계속 작성", f"오류 발생: {str(e)}")]
+
+def stop_generation():
+    stop_event.set()
+    return "생성이 중단되었습니다."
 
-# Gradio 인터페이스 수정
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot()
     msg = gr.Textbox(label="메시지 입력")
@@ -116,7 +106,7 @@ with gr.Blocks() as demo:
     send.click(respond, inputs=[msg, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
     msg.submit(respond, inputs=[msg, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
     continue_btn.click(continue_writing, 
-                       inputs=[lambda: chatbot[-1][1] if chatbot else "", system_message, max_tokens, temperature, top_p, model], 
+                       inputs=[chatbot, system_message, max_tokens, temperature, top_p, model], 
                        outputs=[chatbot])
     stop.click(stop_generation, outputs=[msg])
     clear.click(lambda: None, outputs=[chatbot])