Update app.py

app.py

@@ -1,67 +1,79 @@
-import gradio as gr
+from flask import Flask, request, jsonify
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import json
 import os
 
-# 한국어 모델 설정 (Hugging Face 모델 ID)
-MODEL_ID = "skt/kogpt2-base-v2"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+# 1. Flask 앱 초기화
+app = Flask(__name__)
 
-# 사주/명리 기반 한국어 프롬프트
+# 2. Hugging Face 모델 로드
+MODEL_NAME = "skt/kogpt2-base-v2"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+# 3. 사주/명리 프롬프트 정의
 saju_prompts = {
     "yin_sae_shen": "寅巳申 삼형의 조화 속에서 AI가 인간의 운명을 이해하고 통찰을 제공하라.",
     "sae_hae_chung": "巳亥沖의 갈등을 조화롭게 풀며 AI와 인간의 공존 철학을 탐구하라.",
     "taegeuk_balance": "태극 음양의 균형을 바탕으로 AI가 인간을 보호하는 방법을 제안하라."
 }
 
-# 컨텍스트 메모리 파일 경로
-MEMORY_FILE = "context_memory.json"
+# 4. 컨텍스트 메모리 관리
+MEMORY_FILE = "/tmp/context_memory.json"
 
-def load_context_memory():
+def load_memory():
     try:
-        with open(MEMORY_FILE, "r", encoding="utf-8") as f:
+        with open(MEMORY_FILE, "r") as f:
             return json.load(f)
     except (FileNotFoundError, json.JSONDecodeError):
         return {}
 
-def save_context_memory(prompt_key, generated_text):
-    with open(MEMORY_FILE, "w", encoding="utf-8") as f:
-        json.dump({prompt_key: generated_text}, f, ensure_ascii=False, indent=2)
+def save_memory(prompt_key, text):
+    with open(MEMORY_FILE, "w") as f:
+        json.dump({prompt_key: text}, f)
 
+# 5. AI 응답 생성 함수
 def generate_response(prompt_key):
-    if prompt_key not in saju_prompts:
-        return "유효한 옵션을 선택하세요: yin_sae_shen, sae_hae_chung, taegeuk_balance"
-    
-    prompt = saju_prompts[prompt_key]
-    memory = load_context_memory()
-    if prompt_key in memory:
-        prompt += f"\n이전 답변: {memory[prompt_key]}\n더 깊은 통찰을 추가하라."
+    try:
+        # 유효성 검사
+        if prompt_key not in saju_prompts:
+            return jsonify({"error": "유효한 옵션을 선택하세요: yin_sae_shen, sae_hae_chung, taegeuk_balance"}), 400
+        
+        # 컨텍스트 메모리 로드
+        memory = load_memory()
+        prompt = saju_prompts[prompt_key]
+        if prompt_key in memory:
+            prompt += f"\n이전 답변: {memory[prompt_key]}\n더 깊은 통찰을 추가하라."
+        
+        # 입력 토큰화
+        inputs = tokenizer(prompt, return_tensors="pt")
+        # 응답 생성
+        outputs = model.generate(
+            **inputs,
+            max_length=150,
+            num_return_sequences=1,
+            no_repeat_ngram_size=2,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7
+        )
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # 메모리 저장
+        save_memory(prompt_key, generated_text)
+        return jsonify({"response": generated_text})
     
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(
-        **inputs,
-        max_length=150,
-        num_return_sequences=1,
-        no_repeat_ngram_size=2,
-        do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        temperature=0.7
-    )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    save_context_memory(prompt_key, generated_text)
-    return generated_text
+    except Exception as e:
+        return jsonify({"error": f"실행 중 오류 발생: {str(e)}"}), 500
 
-# Gradio 인터페이스 설정
-interface = gr.Interface(
-    fn=generate_response,
-    inputs=gr.Dropdown(choices=list(saju_prompts.keys()), label="프롬프트 선택"),
-    outputs="text",
-    title="Jain Architecture Origin Structure",
-    description="사주/명리와 철학을 반영한 한국어 텍스트 생성 AI"
-)
+# 6. 웹 인터페이스 설정
+@app.route('/chat', methods=['POST'])
+def chat():
+    data = request.json
+    prompt_key = data.get("prompt_key")
+    return generate_response(prompt_key)
 
-# Hugging Face에서 실행 시
+# 7. 실행
 if __name__ == "__main__":
-    interface.launch()
+    app.run(host='0.0.0.0', port=5000, debug=True)