Update app.py

app.py

@@ -2,37 +2,52 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from retriever.vectordb_rerank import search_documents  # 🧠 RAG 검색기 불러오기
 
 model_name = "dasomaru/gemma-3-4bit-it-demo"
 
-# 🚀 tokenizer는 CPU에서도 미리 불러올 수 있음
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-# 🚀 model은 CPU로만 먼저 올림 (GPU 아직 없음)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,  # 4bit model이니까
-    trust_remote_code=True,
-)
-
 @spaces.GPU(duration=300)
-def generate_response(prompt):
-    # 모델 및 토크나이저 로딩은 함수 내부에서 수행
-    tokenizer = AutoTokenizer.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
-    model = AutoModelForCausalLM.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
-    model.to("cuda")
-
+def generate_response(query):
+    # 모델 및 토크나이저 로드 (ZeroGPU 예약 후)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        trust_remote_code=True,
+    ).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("cuda")
-    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7,
-    top_p=0.9,
-    top_k=50,
-    do_sample=True,)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
+        top_k=50,
+        do_sample=True,
+    )
+
+    # 4. 결과 반환
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+# Gradio 앱
 demo = gr.Interface(fn=generate_response, inputs="text", outputs="text")
 demo.launch()
 
-
-
 # zero = torch.Tensor([0]).cuda()
 # print(zero.device) # <-- 'cpu' 🤔