Update app.py

app.py

@@ -11,65 +11,75 @@ import fitz  # PyMuPDF
 # 환경 변수에서 Hugging Face 토큰 가져오기
 token = os.environ.get("HF_TOKEN")
 
-# 임베딩 모델 로드
-ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
-
-# Lazy Loading PDF 텍스트와 임베딩
+# Lazy Loading 변수
+ST = None
+model = None
+tokenizer = None
 law_sentences = None
 law_embeddings = None
 index = None
-
+data = None
+
+# 임베딩 모델 Lazy Loading
+def load_embedding_model():
+    global ST
+    if ST is None:
+        ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
+    return ST
+
+# LLaMA 모델 및 토크나이저 Lazy Loading
+def load_model():
+    global model, tokenizer
+    if model is None or tokenizer is None:
+        model_id = "google/gemma-2-2b-it"
+        tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            token=token
+        )
+    return model, tokenizer
+
+# PDF에서 텍스트 추출 및 임베딩 Lazy Loading
 def load_law_data():
     global law_sentences, law_embeddings, index
     if law_sentences is None or law_embeddings is None or index is None:
-        # PDF에서 텍스트 추출
         pdf_path = "laws.pdf"  # 여기에 실제 PDF 경로를 입력하세요.
         doc = fitz.open(pdf_path)
         law_text = ""
         for page in doc:
             law_text += page.get_text()
         
-        # 텍스트를 문장 단위로 나누고 임베딩
         law_sentences = law_text.split('\n')  # PDF 구조에 따라 분할을 조정
-        law_embeddings = ST.encode(law_sentences)
+        law_embeddings = load_embedding_model().encode(law_sentences)
 
         # FAISS 인덱스 생성 및 임베딩 추가
         index = faiss.IndexFlatL2(law_embeddings.shape[1])
         index.add(law_embeddings)
 
-# Hugging Face에서 법률 상담 데이터셋 로드
-dataset = load_dataset("jihye-moon/LawQA-Ko")
-data = dataset["train"]
-
-# 질문 컬럼을 임베딩하여 새로운 컬럼에 추가
-data = data.map(lambda x: {"question_embedding": ST.encode(x["question"])}, batched=True)
-data.add_faiss_index(column="question_embedding")
-
-# LLaMA 모델 설정 (양자화 없이)
-model_id = "google/gemma-2-2b-it"
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,  # 양자화 없이 bfloat16 사용
-    device_map="auto",
-    token=token
-)
-
-SYS_PROMPT = """You are an assistant for answering legal questions.
-... (이하 생략, 기존 SYS_PROMPT 그대로 유지) ...
-"""
+# Hugging Face에서 법률 상담 데이터셋 로드 (Lazy Loading)
+def load_dataset_data():
+    global data
+    if data is None:
+        dataset = load_dataset("jihye-moon/LawQA-Ko")
+        data = dataset["train"]
+        data = data.map(lambda x: {"question_embedding": load_embedding_model().encode(x["question"])}, batched=True)
+        data.add_faiss_index(column="question_embedding")
+    return data
 
 # 법률 문서 검색 함수
 def search_law(query, k=5):
     load_law_data()  # PDF 텍스트와 임베딩 Lazy Loading
-    query_embedding = ST.encode([query])
+    query_embedding = load_embedding_model().encode([query])
     D, I = index.search(query_embedding, k)
     return [(law_sentences[i], D[0][idx]) for idx, i in enumerate(I[0])]
 
 # 법률 상담 데이터 검색 함수
 def search_qa(query, k=3):
-    scores, retrieved_examples = data.get_nearest_examples(
-        "question_embedding", ST.encode(query), k=k
+    dataset_data = load_dataset_data()
+    scores, retrieved_examples = dataset_data.get_nearest_examples(
+        "question_embedding", load_embedding_model().encode(query), k=k
     )
     return [retrieved_examples["answer"][i] for i in range(k)]
 
@@ -92,6 +102,8 @@ def talk(prompt, history):
     formatted_prompt = format_prompt(prompt, retrieved_law_docs, qa_results)
     formatted_prompt = formatted_prompt[:2000]  # GPU 메모리 부족을 피하기 위해 프롬프트 제한
 
+    model, tokenizer = load_model()
+
     messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]
 
     # 모델에게 생성 지시