Update app.py

app.py

@@ -1,14 +1,8 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-import os
 import random
-import logging
-
-# 로깅 설정
-logging.basicConfig(filename='language_model_playground.log', level=logging.DEBUG, 
-                    format='%(asctime)s - %(levelname)s - %(message)s')
+import os
+from huggingface_hub import InferenceClient
 
-# 모델 목록
 MODELS = {
     "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
     "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
@@ -21,62 +15,38 @@ MODELS = {
     "Aya-23-35B": "CohereForAI/aya-23-35B"
 }
 
-# HuggingFace 토큰 설정
-hf_token = os.getenv("HF_TOKEN")
-if not hf_token:
-    raise ValueError("HF_TOKEN 환경 변수가 설정되지 않았습니다.")
+def create_client(model_name):
+    return InferenceClient(model_name, token=os.getenv("HF_TOKEN"))
 
-def call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, model):
-    client = InferenceClient(model=model, token=hf_token)
-    combined_prompt = f"{prompt}\n\n참고 텍스트:\n{reference_text}"
+def call_api(model, content, system_message, max_tokens, temperature, top_p):
+    client = create_client(MODELS[model])
+    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
     random_seed = random.randint(0, 1000000)
-    
-    try:
-        response = client.text_generation(
-            combined_prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            seed=random_seed
-        )
-        return response
-    except Exception as e:
-        logging.error(f"HuggingFace API 호출 중 오류 발생: {str(e)}")
-        return f"응답 생성 중 오류 발생: {str(e)}. 나중에 다시 시도해 주세요."
+    response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, seed=random_seed)
+    return response.choices[0].message.content
 
-def generate_response(prompt, reference_text, max_tokens, temperature, top_p, model):
-    response = call_hf_api(prompt, reference_text, max_tokens, temperature, top_p, MODELS[model])
-    response_html = f"""
-    <h3>생성된 응답:</h3>
-    <div style='max-height: 500px; overflow-y: auto; white-space: pre-wrap; word-wrap: break-word;'>
-    {response}
-    </div>
-    """
-    return response_html
+def generate_text(model, user_message, system_message, max_tokens, temperature, top_p):
+    return call_api(model, user_message, system_message, max_tokens, temperature, top_p)
 
-# Gradio 인터페이스 설정
-with gr.Blocks() as demo:
-    gr.Markdown("## 언어 모델 프롬프트 플레이그라운드")
+title = "AI 텍스트 생성기"
 
-    with gr.Column():
-        model_radio = gr.Radio(choices=list(MODELS.keys()), value="Zephyr 7B Beta", label="언어 모델 선택")
-        prompt_input = gr.Textbox(label="프롬프트 입력", lines=5)
-        reference_text_input = gr.Textbox(label="참고 텍스트 입력", lines=5)
-        
-        with gr.Row():
-            max_tokens_slider = gr.Slider(minimum=0, maximum=5000, value=2000, step=100, label="최대 토큰 수")
-            temperature_slider = gr.Slider(minimum=0, maximum=1, value=0.75, step=0.05, label="온도")
-            top_p_slider = gr.Slider(minimum=0, maximum=1, value=0.95, step=0.05, label="Top P")
-        
-        generate_button = gr.Button("응답 생성")
-        response_output = gr.HTML(label="생성된 응답")
-
-    # 버튼 클릭 시 응답 생성
-    generate_button.click(
-        generate_response,
-        inputs=[prompt_input, reference_text_input, max_tokens_slider, temperature_slider, top_p_slider, model_radio],
-        outputs=response_output
-    )
+with gr.Blocks() as demo:
+    gr.Markdown(f"# {title}")
+    
+    model = gr.Radio(choices=list(MODELS.keys()), label="언어 모델 선택", value="Zephyr 7B Beta")
+    user_message = gr.Textbox(label="사용자 메시지", lines=5)
+    system_message = gr.Textbox(label="시스템 메시지 (프롬프트)", lines=10)
+    
+    with gr.Accordion("고급 설정", open=False):
+        max_tokens = gr.Slider(label="Max Tokens", minimum=0, maximum=4000, value=500, step=100)
+        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.75, step=0.05)
+        top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
+    
+    generate_btn = gr.Button("텍스트 생성하기")
+    output = gr.Textbox(label="생성된 텍스트", lines=10)
+    
+    generate_btn.click(fn=generate_text, 
+                       inputs=[model, user_message, system_message, max_tokens, temperature, top_p], 
+                       outputs=[output])
 
-# 인터페이스 실행
-demo.launch(share=True)
+demo.launch()