app.py

import gradio as gr
from transformers import pipeline
import threading

chat_model = None
loading_done = False
status_text = "⏳ 모델 로딩 중입니다..."

# 백그라운드에서 모델 로딩
def load_model_bg():
    global chat_model, loading_done, status_text
    chat_model = pipeline("text-generation", model="beomi/KoAlpaca-Polyglot-5.8B")
    loading_done = True
    status_text = "✅ 모델 로딩 완료! 문장을 입력해보세요."

# 대답 함수
def chat_with_model(message, history):
    global chat_model, loading_done
    if not loading_done:
        return history + [[message, "⚠️ 모델이 아직 로딩 중입니다. 잠시만 기다려 주세요."]]
    
    prompt = f"""
    ### Instruction: 다음 문장을 분석해서 무례하거나 공격적인 표현이 있는지 판단하고, 있다면 더 예의 있게 고쳐줘.
    
    ### Input:
    {message}
    
    ### Response:
    """
    response = chat_model(prompt, max_new_tokens=200)[0]['generated_text']
    response_only = response[len(prompt):].strip()
    return history + [[message, response_only]]

# 상태 텍스트 반환 함수 (매번 새로 읽어옴)
def get_status():
    return status_text

# 백그라운드에서 모델 로딩 시작
threading.Thread(target=load_model_bg).start()

# Gradio 앱
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="문장을 입력하세요", placeholder="예: 너 정말 왜 그렇게 말해?")
    status = gr.Markdown(get_status)

    def respond_and_clear(user_input, history):
        updated_history = chat_with_model(user_input, history)
        return "", updated_history, get_status()

    msg.submit(respond_and_clear, [msg, chatbot], [msg, chatbot, status])

demo.launch()