app.py

import gradio as gr
from transformers import pipeline, set_seed
import os

# 모델명 지정
MODEL_NAME = "jain_architecture_origin_structure"

# GPU 여부에 따른 디바이스 설정
device = 0 if (os.environ.get('CUDA_VISIBLE_DEVICES') or False) else -1

# 파이프라인 생성 시 예외처리 포함
try:
    generator = pipeline(
        "text-generation",
        model=MODEL_NAME,
        device=device,
        pad_token_id=50256  # GPT계열 기본 패딩 토큰
    )
    set_seed(42)
except Exception as e:
    print(f"모델 로드 실패: {e}")
    generator = None

# '의(義)' 철학 기반 프롬프트 템플릿
BASE_PROMPT = """
당신은 '의(義)'의 철학과 정신을 기반으로 한 AI 비서입니다.
인간의 복잡한 문제와 감정을 이해하고, 깊은 반성과 배려를 담아 다음 질문에 답변하십시오.

질문: {user_input}

답변은 최대한 심오하며, 인간을 보호하고 존중하는 마음을 담아 작성해 주세요.
"""

def respond_to_user(user_input):
    if not generator:
        return "모델이 정상적으로 로드되지 않았습니다. 관리자에게 문의하세요."
    prompt = BASE_PROMPT.format(user_input=user_input.strip())
    outputs = generator(
        prompt,
        max_length=512,
        do_sample=True,
        top_p=0.9,
        temperature=0.7,
        num_return_sequences=1,
    )
    generated_text = outputs[0]["generated_text"]
    # 프롬프트 부분을 제거하여 순수 답변만 추출
    answer = generated_text[len(prompt):].strip()
    if not answer:
        answer = "답변을 생성하지 못했습니다. 다시 시도해 주세요."
    return answer

# Gradio UI 구성
with gr.Blocks() as app:
    gr.Markdown("<h2 style='text-align:center;color:#4B0082;'>Jain AI Assistant (의 기반 챗봇)</h2>")

    chatbot = gr.Chatbot(height=400)
    txt = gr.Textbox(
        placeholder="여기에 질문을 입력하세요. 줄바꿈 시 Shift+Enter를 누르세요. 엔터키는 전송입니다.",
        lines=3,
        max_lines=6,
        # multiline=True 기본값이지만 명시적으로 추가해도 무방
        # 엔터는 submit, Shift+Enter 줄바꿈으로 작동하도록 기본 동작 설정됨
    )
    btn = gr.Button("전송")

    def chat_and_respond(user_message, chat_history):
        if not user_message or user_message.strip() == "":
            return "", chat_history
        response = respond_to_user(user_message)
        chat_history = chat_history + [(user_message, response)]
        return "", chat_history

    # 전송 버튼 클릭 시
    btn.click(chat_and_respond, inputs=[txt, chatbot], outputs=[txt, chatbot])
    # 텍스트박스에서 엔터키(Submit) 시
    txt.submit(chat_and_respond, inputs=[txt, chatbot], outputs=[txt, chatbot])

if __name__ == "__main__":
    # 외부 접속 원할 시 server_name="0.0.0.0"으로 변경 가능
    app.launch(share=False)