Update app.py

app.py

@@ -7,6 +7,7 @@ from threading import Thread
 from sentence_transformers import SentenceTransformer
 import faiss
 import fitz  # PyMuPDF
+import spaces
 
 # 환경 변수에서 Hugging Face 토큰 가져오기
 token = os.environ.get("HF_TOKEN")
@@ -21,6 +22,7 @@ index = None
 data = None
 
 # 임베딩 모델 Lazy Loading
+@spaces.GPU
 def load_embedding_model():
     global ST
     if ST is None:
@@ -28,6 +30,7 @@ def load_embedding_model():
     return ST
 
 # LLaMA 모델 및 토크나이저 Lazy Loading
+@spaces.GPU
 def load_model():
     global model, tokenizer
     if model is None or tokenizer is None:
@@ -42,6 +45,7 @@ def load_model():
     return model, tokenizer
 
 # PDF에서 텍스트 추출 및 임베딩 Lazy Loading
+@spaces.GPU
 def load_law_data():
     global law_sentences, law_embeddings, index
     if law_sentences is None or law_embeddings is None or index is None:
@@ -59,6 +63,7 @@ def load_law_data():
         index.add(law_embeddings)
 
 # Hugging Face에서 법률 상담 데이터셋 로드 (Lazy Loading)
+@spaces.GPU
 def load_dataset_data():
     global data
     if data is None:
@@ -68,7 +73,9 @@ def load_dataset_data():
         data.add_faiss_index(column="question_embedding")
     return data
 
+
 # 법률 문서 검색 함수
+@spaces.GPU
 def search_law(query, k=5):
     load_law_data()  # PDF 텍스트와 임베딩 Lazy Loading
     query_embedding = load_embedding_model().encode([query])
@@ -76,6 +83,7 @@ def search_law(query, k=5):
     return [(law_sentences[i], D[0][idx]) for idx, i in enumerate(I[0])]
 
 # 법률 상담 데이터 검색 함수
+@spaces.GPU
 def search_qa(query, k=3):
     dataset_data = load_dataset_data()
     scores, retrieved_examples = dataset_data.get_nearest_examples(
@@ -84,6 +92,7 @@ def search_qa(query, k=3):
     return [retrieved_examples["answer"][i] for i in range(k)]
 
 # 최종 프롬프트 생성
+@spaces.GPU
 def format_prompt(prompt, law_docs, qa_docs):
     PROMPT = f"Question: {prompt}\n\nLegal Context:\n"
     for doc in law_docs:
@@ -94,6 +103,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)
@@ -157,4 +167,5 @@ demo = gr.ChatInterface(
 )
 
 # Gradio 데모 실행
+@spaces.GPU
 demo.launch(debug=True, server_port=7860)