Update app.py

app.py

@@ -1,24 +1,19 @@
-from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import InferenceClient
 import json
 import os
 
-# 1. Flask 앱 초기화
-app = Flask(__name__)
+# Hugging Face 모델 설정
+MODEL_ID = "skt/kogpt2-base-v2"  # 한국어 모델 ID
+CLIENT = InferenceClient(model=MODEL_ID)
 
-# 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가 인간을 보호하는 방법을 제안하라."
 }
 
-# 4. 컨텍스트 메모리 관리
+# 컨텍스트 메모리 경로
 MEMORY_FILE = "/tmp/context_memory.json"
 
 def load_memory():
@@ -32,12 +27,13 @@ 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):
+def handle_request(request_data):
     try:
+        prompt_key = request_data.get("prompt_key")
+        
         # 유효성 검사
         if prompt_key not in saju_prompts:
-            return jsonify({"error": "유효한 옵션을 선택하세요: yin_sae_shen, sae_hae_chung, taegeuk_balance"}), 400
+            return {"error": "유효한 옵션을 선택하세요: yin_sae_shen, sae_hae_chung, taegeuk_balance"}
         
         # 컨텍스트 메모리 로드
         memory = load_memory()
@@ -45,35 +41,35 @@ def generate_response(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,
+        # Hugging Face API 호출
+        response = CLIENT.chat(
+            model=MODEL_ID,
+            messages=[
+                {"role": "system", "content": prompt},
+                {"role": "user", "content": "분석을 시작해 주세요."}
+            ],
+            max_tokens=400,
             temperature=0.7
         )
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # 메모리 저장
-        save_memory(prompt_key, generated_text)
-        return jsonify({"response": generated_text})
+        # 결과 처리 및 메모리 저장
+        result = response.choices[0].message.content
+        save_memory(prompt_key, result)
+        return {"response": result}
     
     except Exception as e:
-        return jsonify({"error": f"실행 중 오류 발생: {str(e)}"}), 500
+        return {"error": f"실행 중 오류 발생: {str(e)}"}
 
-# 6. 웹 인터페이스 설정
-@app.route('/chat', methods=['POST'])
-def chat():
-    data = request.json
-    prompt_key = data.get("prompt_key")
-    return generate_response(prompt_key)
-
-# 7. 실행
+# Hugging Face 실행 환경에서의 요청 처리
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=5000, debug=True)
+    # Hugging Face는 `request` 객체를 제공하지 않으므로, 아래와 같이 대체
+    # 실제 환경에서는 `request` 대신 환경 변수나 입력값을 사용해야 함
+    import sys
+    if len(sys.argv) < 2:
+        print("Usage: python app.py <request_data>")
+        sys.exit(1)
+    
+    # 요청 데이터 파싱 (예: {"prompt_key": "yin_sae_shen"})
+    request_data = json.loads(sys.argv[1])
+    result = handle_request(request_data)
+    print(json.dumps(result))