Update app.py

app.py

@@ -11,72 +11,65 @@ import fitz  # PyMuPDF
 # 환경 변수에서 Hugging Face 토큰 가져오기
 token = os.environ.get("HF_TOKEN")
 
-# 임베딩 모델 Lazy Loading
-ST = None
-def load_embedding_model():
-    global ST
-    if ST is None:
-        ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
-    return ST
-
-# LLaMA 모델 및 토크나이저 Lazy Loading
-model = None
-tokenizer = None
-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에서 텍스트 추출
-def extract_text_from_pdf(pdf_path):
-    doc = fitz.open(pdf_path)
-    text = ""
-    for page in doc:
-        text += page.get_text()
-    return text
-
-# 법률 문서 PDF 경로 지정 및 텍스트 추출
-pdf_path = "laws.pdf"  # 여기에 실제 PDF 경로를 입력하세요.
-law_text = extract_text_from_pdf(pdf_path)
-
-# 법률 문서 텍스트를 문장 단위로 나누고 임베딩
-law_sentences = law_text.split('\n')  # PDF 구조에 따라 분할을 조정
-law_embeddings = load_embedding_model().encode(law_sentences)
-
-# FAISS 인덱스 생성 및 임베딩 추가
-index = faiss.IndexFlatL2(law_embeddings.shape[1])
-index.add(law_embeddings)
-
-# Hugging Face에서 법률 상담 데이터셋 로드 (Lazy Loading)
-data = None
-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
+# 임베딩 모델 로드
+ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
+
+# Lazy Loading PDF 텍스트와 임베딩
+law_sentences = None
+law_embeddings = None
+index = None
+
+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)
+
+        # 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 그대로 유지) ...
+"""
 
 # 법률 문서 검색 함수
 def search_law(query, k=5):
-    query_embedding = load_embedding_model().encode([query])
+    load_law_data()  # PDF 텍스트와 임베딩 Lazy Loading
+    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])]
 
 # 법률 상담 데이터 검색 함수
 def search_qa(query, k=3):
-    dataset_data = load_dataset_data()
-    scores, retrieved_examples = dataset_data.get_nearest_examples(
-        "question_embedding", load_embedding_model().encode(query), k=k
+    scores, retrieved_examples = data.get_nearest_examples(
+        "question_embedding", ST.encode(query), k=k
     )
     return [retrieved_examples["answer"][i] for i in range(k)]
 
@@ -99,8 +92,6 @@ 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}]
 
     # 모델에게 생성 지시
@@ -155,4 +146,3 @@ demo = gr.ChatInterface(
 
 # Gradio 데모 실행
 demo.launch(debug=True, server_port=7860)
-