Update app.py

app.py

@@ -1,16 +1,12 @@
 import gradio as gr
 from datasets import load_dataset
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
 from threading import Thread
 from sentence_transformers import SentenceTransformer
 import faiss
 import fitz  # PyMuPDF
-
-# 환경 변수에서 Hugging Face 토큰 가져오기
-token = os.environ.get("HF_TOKEN")
-
+import os
 
 # 임베딩 모델 로드
 ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
@@ -23,17 +19,36 @@ def extract_text_from_pdf(pdf_path):
         text += page.get_text()
     return text
 
-# 법률 문서 PDF 경로 지정 및 텍스트 추출
-pdf_path = "laws.pdf"  # 여기에 실제 PDF 경로를 입력하세요.
-law_text = extract_text_from_pdf(pdf_path)
+# 기본 제공 PDF 파일 경로
+default_pdf_path = "laws.pdf"
+
+# FAISS 인덱스 초기화
+index = None
+law_sentences = []
+
+# 기본 제공 PDF 파일 처리 함수
+def process_default_pdf():
+    global index, law_sentences
 
-# 법률 문서 텍스트를 문장 단위로 나누고 임베딩
-law_sentences = law_text.split('\n')  # Adjust splitting based on your PDF structure
-law_embeddings = ST.encode(law_sentences)
+    # PDF에서 텍스트 추출
+    law_text = extract_text_from_pdf(default_pdf_path)
 
-# FAISS 인덱스 생성 및 임베딩 추가
-index = faiss.IndexFlatL2(law_embeddings.shape[1])
-index.add(law_embeddings)
+    # 문장을 나누고 임베딩 생성
+    law_sentences = law_text.split('\n')
+    law_embeddings = ST.encode(law_sentences)
+
+    # FAISS 인덱스 생성 및 임베딩 추가
+    index = faiss.IndexFlatL2(law_embeddings.shape[1])
+    index.add(law_embeddings)
+
+# 처음에 기본 PDF 파일 처리
+process_default_pdf()
+
+# 법률 문서 검색 함수
+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])]
 
 # Hugging Face에서 법률 상담 데이터셋 로드
 dataset = load_dataset("jihye-moon/LawQA-Ko")
@@ -44,29 +59,21 @@ data = data.map(lambda x: {"question_embedding": ST.encode(x["question"])}, batc
 data.add_faiss_index(column="question_embedding")
 
 # LLaMA 모델 설정
-model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
-)
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+model_id = "google/gemma-2-2b-it"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     torch_dtype=torch.bfloat16,
     device_map="auto",
-    quantization_config=bnb_config,
-    token=token
 )
 
 SYS_PROMPT = """You are an assistant for answering legal questions.
 You are given the extracted parts of legal documents and a question. Provide a conversational answer.
 If you don't know the answer, just say "I do not know." Don't make up an answer.
-you must answer korean."""
-
-# 법률 문서 검색 함수
-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])]
+you must answer korean.
+You're a LAWEYE legal advisor bot. Your job is to provide korean legal assistance by asking questions to korean speaker, then offering advice or guidance based on the information and law provisions provided. Make sure you only respond with one question at a time.
+...
+"""
 
 # 법률 상담 데이터 검색 함수
 def search_qa(query, k=3):
@@ -97,11 +104,7 @@ def talk(prompt, history):
     messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]
 
     # 모델에게 생성 지시
-    input_ids = tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to(model.device)
+    input_ids = tokenizer(messages, return_tensors="pt").to(model.device).input_ids
 
     streamer = TextIteratorStreamer(
         tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
@@ -147,4 +150,4 @@ demo = gr.ChatInterface(
 )
 
 # Gradio 데모 실행
-demo.launch(debug=True)
+demo.launch(debug=True)