Update app.py

app.py

@@ -1,3 +1,4 @@
+import space
 import gradio as gr
 from datasets import load_dataset
 import os
@@ -8,6 +9,7 @@ from sentence_transformers import SentenceTransformer
 import faiss
 import fitz  # PyMuPDF
 
+
 # 환경 변수에서 Hugging Face 토큰 가져오기
 token = os.environ.get("HF_TOKEN")
 
@@ -62,13 +64,15 @@ You are given the extracted parts of legal documents and a question. Provide a c
 If you don't know the answer, just say "I do not know." Don't make up an answer.
 you must answer korean."""
 
-# 법률 문서 검색 함수
+# 법률 문서 검색 
+@spaces.Gpu
 def search_law(query, k=5):
     query_embedding = ST.encode([query])
     D, I = index.search(query_embedding, k)
     return [(law_sentences[i], D[0][idx]) for idx, i in enumerate(I[0])]
 
 # 법률 상담 데이터 검색 함수
+@spaces.Gpu
 def search_qa(query, k=3):
     scores, retrieved_examples = data.get_nearest_examples(
         "question_embedding", ST.encode(query), k=k
@@ -86,6 +90,7 @@ def format_prompt(prompt, law_docs, qa_docs):
     return PROMPT
 
 # 챗봇 응답 함수
+@spaces.Gpu
 def talk(prompt, history):
     law_results = search_law(prompt, k=3)
     qa_results = search_qa(prompt, k=3)