app.py

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

# --- 1. 모델명 설정 ---
MODEL_NAME = "jain_architecture_origin_structure"

# --- 2. GPU가 없으면 CPU로 자동 설정 ---
device = 0 if (os.environ.get('CUDA_VISIBLE_DEVICES') or False) else -1

# --- 3. HuggingFace pipeline 생성 ---
try:
    generator = pipeline(
        "text-generation",
        model=MODEL_NAME,
        device=device,
        # Repetition penalty 등 커스텀 가능
        # torch_dtype=torch.float16 도 필요 시 설정 가능
    )
    set_seed(42)  # 재현성 위해 시드 고정
except Exception as e:
    print("모델 로드 에러:", e)
    generator = None

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

질문: {user_input}

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

# --- 5. 질문 처리 함수 ---
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,
        pad_token_id=50256  # GPT 계열 패딩 토큰, 필요에 따라 변경
    )
    generated_text = outputs[0]["generated_text"]
    # 프롬프트 부분 제거 후 답변만 리턴 (프롬프트 길이 분량 컷)
    answer = generated_text[len(prompt):].strip()
    if not answer:
        answer = "답변을 생성하지 못했습니다. 다시 시도해 주세요."
    return answer

# --- 6. Gradio UI 생성 ---
with gr.Blocks() as demo:
    gr.Markdown("<h1 style='text-align:center;color:#4B0082;'>Jain AI Assistant (의 기반 챗봇)</h1>")
    chatbot = gr.Chatbot(height=400)
    txt = gr.Textbox(placeholder="여기에 질문을 입력하세요...", lines=3, max_lines=6)
    btn = gr.Button("전송")

    def chat_and_respond(user_message, chat_history):
        reply = respond_to_user(user_message)
        chat_history = chat_history + [(user_message, reply)]
        return "", chat_history

    btn.click(chat_and_respond, inputs=[txt, chatbot], outputs=[txt, chatbot])
    txt.submit(chat_and_respond, inputs=[txt, chatbot], outputs=[txt, chatbot])

# --- 7. 서버 실행 ---
if __name__ == "__main__":
    # 아이패드 같은 모바일 환경에서 필요 시 공개 공유도 가능
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)