Update app.py

app.py

@@ -8,7 +8,7 @@ from services.rag_pipeline import rag_pipeline
 model_name = "dasomaru/gemma-3-4bit-it-demo"
 
 
-# 🚀 tokenizer는 CPU에서도 미리 불러올 수 있음
+# 1. 모델/토크나이저 1회 로딩
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 # 🚀 model은 CPU로만 먼저 올림 (GPU 아직 없음)
 model = AutoModelForCausalLM.from_pretrained(
@@ -17,65 +17,37 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
 )
 
-# v0
+# 2. 캐시 관리
+search_cache = {}
+
 @spaces.GPU(duration=300)
-def generate_response(query):
-    # 🚀 generate_response 함수 안에서 매번 로드
-    # tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    # model = AutoModelForCausalLM.from_pretrained(
-    #     model_name,
-    #     torch_dtype=torch.float16,
-    #     device_map="auto",  # ✅ 중요: 자동으로 GPU 할당
-    #     trust_remote_code=True,
-    # )
+def generate_response(query: str):
     tokenizer = AutoTokenizer.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
     model = AutoModelForCausalLM.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
     model.to("cuda")    
 
-    # 1. 검색
-    top_k = 5
-    retrieved_docs = search_documents(query, top_k=top_k)
-
-    # 2. 프롬프트 조립
-    prompt = (
-        "당신은 공인중개사 시험 문제 출제 전문가입니다.\n\n"
-        "다음은 기출 문제 및 관련 법령 정보입니다:\n"
-    )
-    for idx, doc in enumerate(retrieved_docs, 1):
-        prompt += f"- {doc}\n"
-    prompt += f"\n이 정보를 참고하여 사용자의 요청에 답변해 주세요.\n\n"
-    prompt += f"[질문]\n{query}\n\n[답변]\n"
-
-    # 3. 답변 생성
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)  # ✅ model.device
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        top_k=50,
-        do_sample=True,
-    )
-
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# v1
-@spaces.GPU(duration=300)
-def generate_response_with_pipeline(query):
-    return rag_pipeline(query)
-    
-# v2
-search_cache = {}
-@spaces.GPU(duration=300)
-def search_documents_with_cache(query: str):
     if query in search_cache:
         print(f"⚡ 캐시 사용: '{query}'")
         return search_cache[query]
 
+    # 🔥 rag_pipeline을 호출해서 검색 + 생성
     results = rag_pipeline(query)
+
+    # 결과가 list일 경우 합치기
+    if isinstance(results, list):
+        results = "\n\n".join(results)
+
     search_cache[query] = results
     return results
+    
+# 3. Gradio 인터페이스
+demo = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(lines=2, placeholder="질문을 입력하세요"),
+    outputs="text",
+    title="Law RAG Assistant",
+    description="법령 기반 RAG 파이프라인 테스트",
+)
 
-
-demo = gr.Interface(fn=search_documents_with_cache, inputs="text", outputs="text")
+# demo.launch(server_name="0.0.0.0", server_port=7860)  # 🚀 API 배포 준비 가능
 demo.launch()