Update app.py

app.py

@@ -1,118 +1,62 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-import os
-from threading import Event
-import re
 
-hf_token = os.getenv("HF_TOKEN")
-stop_event = Event()
-
-# 모델 목록 정의
-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"
+MODELS = {
+    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
+    "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
+    "Meta Llama 3.1 8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
+    "Mixtral 8x7B": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "Cohere Command R+": "CohereForAI/c4ai-command-r-plus",
 }
 
-# Inference 클라이언트 반환
-def get_client(model):
-    return InferenceClient(model=model, token=hf_token)
-
-# 응답 필터링 함수
-def filter_response(response):
-    # 질문 형태 제거
-    response = re.sub(r'\?|질문|물어보|궁금', '', response)
-    
-    # 추가 대화나 확장을 제안하는 문구 제거
-    response = re.sub(r'더 자세히 알고 싶으시면|추가로 궁금한 점이 있으시면|다른 질문이 있으신가요?', '', response)
-    
-    # 불필요한 공백 제거 및 문장 끝 정리
-    response = re.sub(r'\s+', ' ', response).strip()
-    response = re.sub(r'[.!]\s*$', '', response) + '.'
-    
-    return response
-
-# 응답 생성 함수 (스트리밍 방식, 자문자답 방지)
-def respond(prompt, system_message, max_tokens, temperature, top_p, selected_model):
-    stop_event.clear()
-    client = get_client(selected_model)
-    
-    # 프롬프트 설정
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": prompt + "\n절대로 추가 질문을 하거나 대화를 확장하지 마세요. 오직 위의 내용에 대해서만 간결하게 답변하세요."}
-    ]
-    
-    try:
-        full_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,
-            stop_sequences=["Human:", "User:", "System:", "Assistant:"]  # 최대 4개의 stop 시퀀스만 사용
-        ):
-            if stop_event.is_set():
-                break
-            if chunk:
-                full_response += chunk
-                filtered_response = filter_response(full_response)
-                yield [(prompt, filtered_response.strip())]
-    
-    except Exception as e:
-        yield [(prompt, f"오류 발생: {str(e)}")]
-
-# 응답 중단 함수
-def stop_generation():
-    stop_event.set()
-
-# Gradio UI 구성
-with gr.Blocks() as demo:
-    gr.Markdown("# 프롬프트 최적화 Playground")
-    
-    gr.Markdown("""
-    **주의사항:**
-    - '전송' 버튼을 클릭하거나 입력 필드에서 Shift+Enter를 눌러 메시지를 전송할 수 있습니다.
-    - Enter 키는 줄바꿈으로 작동합니다.
-    - 입력한 내용에 대해서만 응답하도록 설정되어 있으며, 모델이 추가 질문을 만들거나 입력을 확장하지 않습니다.
-    """)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            with gr.Accordion("모델 설정", open=True):
-                model = gr.Radio(list(models.keys()), value=list(models.keys())[0], label="언어 모델 선택", info="사용할 언어 모델을 선택하세요")
-                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 (핵 샘플링)")
-            
-            system_message = gr.Textbox(
-                value="당신은 정확하고 간결한 응답만을 제공하는 AI 어시스턴트입니다. 어떤 경우에도 추가 질문을 하거나 대화를 확장하지 마세요. 오직 주어진 입력에 대해서만 답변하세요.",
-                label="시스템 메시지",
-                lines=5
-            )
-        
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(height=400, label="대화 결과")
-            prompt = gr.Textbox(
-                label="내용 입력", 
-                lines=3,
-                placeholder="메시지를 입력하세요. 전송 버튼을 클릭하거나 Shift+Enter를 눌러 전송합니다."
-            )
-            
-            with gr.Row():
-                send = gr.Button("전송")
-                stop = gr.Button("🛑 생성 중단")
-                clear = gr.Button("🗑️ 대화 내역 지우기")
-
-    # Event handlers
-    send.click(respond, inputs=[prompt, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
-    prompt.submit(respond, inputs=[prompt, system_message, max_tokens, temperature, top_p, model], outputs=[chatbot])
-    stop.click(stop_generation)
-    clear.click(lambda: None, outputs=[chatbot])
-
-# UI 실행
-demo.launch()
+def get_client(model_name):
+    return InferenceClient(MODELS[model_name])
+
+def respond(
+    message,
+    history: list[tuple[str, str]],
+    model_name,
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
+    client = get_client(model_name)
+    messages = [{"role": "system", "content": system_message}]
+
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
+
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+
+    for message in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = message.choices[0].delta.content
+        response += token
+        yield response
+
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Dropdown(choices=list(MODELS.keys()), label="Language Model", value="Zephyr 7B Beta"),
+        gr.Textbox(value="You are a friendly and helpful AI assistant.", label="System Message"),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+    ],
+    title="Advanced AI Chatbot",
+    description="Chat with different language models and customize your experience!",
+)
+
+if __name__ == "__main__":
+    demo.launch()