Update app.py

app.py

@@ -2,20 +2,21 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from retriever.vectordb_rerank import search_documents  # 🧠 RAG 검색기 불러오기
+from retriever.vectordb import search_documents  # 🧠 RAG 검색기 불러오기
 
 model_name = "dasomaru/gemma-3-4bit-it-demo"
 
 @spaces.GPU(duration=300)
 def generate_response(query):
-    # 모델 및 토크나이저 로드 (ZeroGPU 예약 후)
+    # 🚀 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,
-    ).to("cuda")
-    
+    )
+
     # 1. 검색
     top_k = 5
     retrieved_docs = search_documents(query, top_k=top_k)
@@ -31,7 +32,7 @@ def generate_response(query):
     prompt += f"[질문]\n{query}\n\n[답변]\n"
 
     # 3. 답변 생성
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)  # ✅ model.device
     outputs = model.generate(
         **inputs,
         max_new_tokens=512,
@@ -41,20 +42,7 @@ def generate_response(query):
         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' 🤔
-
-# @spaces.GPU
-# def greet(n):
-#     print(zero.device) # <-- 'cuda:0' 🤗
-#     return f"Hello {zero + n} Tensor"
-
-# demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-# demo.launch()