Update app.py

app.py

@@ -1114,50 +1114,36 @@ def recommend_content_based(user_profile: dict, top_n=5):
 #####################################
 # 5) 챗봇 로직
 #####################################
-HF_API_KEY = os.environ.get("HF_API_KEY", "YOUR_HF_API_KEY")
-API_URL = "https://api-inference.huggingface.co/models/Chanjeans/tfchatbot_2"
-HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}
+tokenizer = AutoTokenizer.from_pretrained("Chanjeans/tfchatbot_2")
+model = AutoModelForCausalLM.from_pretrained("Chanjeans/tfchatbot_2")
+model.eval()
+print("Model loaded successfully.")
 
-def chat_response(user_input, mode="emotion", max_retries=5):
+def chat_response(user_input, mode="emotion"):
     if mode not in ["emotion", "rational"]:
         raise HTTPException(status_code=400, detail="mode는 'emotion' 또는 'rational'이어야 합니다.")
 
     prompt = f"<{mode}><usr>{user_input}</usr><sys>"
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 128,
-            "temperature": 0.7,
-            "top_p": 0.9,
-            "top_k": 50,
-            "repetition_penalty": 1.2,
-            "do_sample": True
-        },
-        "options": {"wait_for_model": True}
-    }
 
-    for attempt in range(max_retries):
-        response = requests.post(API_URL, headers=HEADERS, json=payload)
-        if response.status_code == 200:
-            try:
-                result = response.json()
-                if isinstance(result, list) and "generated_text" in result[0]:
-                    generated_text = result[0]["generated_text"]
-                    return generated_text.replace(prompt, "").strip()
-                else:
-                    return "응답 형식이 예상과 다릅니다."
-            except Exception as e:
-                return f"JSON 파싱 오류: {e}"
-
-        elif response.status_code == 503:
-            # 모델 로딩 중
-            error_info = response.json()
-            estimated_time = error_info.get("estimated_time", 15)
-            time.sleep(min(estimated_time, 15))
-        else:
-            return f"API Error: {response.status_code}, {response.text}"
-
-    return "🚨 모델 로딩이 너무 오래 걸립니다. 잠시 후 다시 시도하세요."
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            temperature=0.7,
+            top_p=0.9,
+            top_k=50,
+            repetition_penalty=1.2,
+            do_sample=True
+        )
+
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # prompt 부분 제거 (불필요한 프롬프트까지 반환되지 않도록)
+    response_text = generated_text.replace(prompt, "").strip()
+
+    return response_text
+
 
 
 #우울분류 모델 추가