Update app.py

app.py

@@ -7,12 +7,11 @@ from threading import Thread
 from sentence_transformers import SentenceTransformer
 import faiss
 import fitz  # PyMuPDF
-import torchvision
+from google.colab import userdata
 
 # 환경 변수에서 Hugging Face 토큰 가져오기
-token = os.environ.get("HF_TOKEN")
-if not token:
-    raise ValueError("Hugging Face token is missing. Please set it in your environment variables.")
+token = userdata.get("HF_TOKEN")
+
 
 # 임베딩 모델 로드
 ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
@@ -30,7 +29,7 @@ pdf_path = "laws.pdf"  # 여기에 실제 PDF 경로를 입력하세요.
 law_text = extract_text_from_pdf(pdf_path)
 
 # 법률 문서 텍스트를 문장 단위로 나누고 임베딩
-law_sentences = law_text.split('\n')
+law_sentences = law_text.split('\n')  # Adjust splitting based on your PDF structure
 law_embeddings = ST.encode(law_sentences)
 
 # FAISS 인덱스 생성 및 임베딩 추가
@@ -46,7 +45,7 @@ data = data.map(lambda x: {"question_embedding": ST.encode(x["question"])}, batc
 data.add_faiss_index(column="question_embedding")
 
 # LLaMA 모델 설정
-model_id = "MLP-KTLim/llama-3-Korean-Bllossom-8B"
+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
 )
@@ -61,7 +60,8 @@ model = AutoModelForCausalLM.from_pretrained(
 
 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."""
+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):
@@ -80,7 +80,7 @@ def search_qa(query, k=3):
 def format_prompt(prompt, law_docs, qa_docs):
     PROMPT = f"Question: {prompt}\n\nLegal Context:\n"
     for doc in law_docs:
-        PROMPT += f"{doc[0]}\n"
+        PROMPT += f"{doc[0]}\n"  # Assuming doc[0] contains the relevant text
     PROMPT += "\nLegal QA:\n"
     for doc in qa_docs:
         PROMPT += f"{doc}\n"
@@ -90,10 +90,11 @@ def format_prompt(prompt, law_docs, qa_docs):
 def talk(prompt, history):
     law_results = search_law(prompt, k=3)
     qa_results = search_qa(prompt, k=3)
-    
+
     retrieved_law_docs = [result[0] for result in law_results]
     formatted_prompt = format_prompt(prompt, retrieved_law_docs, qa_results)
     formatted_prompt = formatted_prompt[:2000]  # GPU 메모리 부족을 피하기 위해 프롬프트 제한
+
     messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]
 
     # 모델에게 생성 지시
@@ -116,6 +117,7 @@ def talk(prompt, history):
         temperature=0.75,
         eos_token_id=tokenizer.eos_token_id,
     )
+
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
 
@@ -126,11 +128,8 @@ def talk(prompt, history):
 
 # Gradio 인터페이스 설정
 TITLE = "Legal RAG Chatbot"
-
-DESCRIPTION = """
-A chatbot that uses Retrieval-Augmented Generation (RAG) for legal consultation.
-This chatbot can search legal documents and previous legal QA pairs to provide answers.
-"""
+DESCRIPTION = """A chatbot that uses Retrieval-Augmented Generation (RAG) for legal consultation.
+This chatbot can search legal documents and previous legal QA pairs to provide answers."""
 
 demo = gr.ChatInterface(
     fn=talk,
@@ -149,5 +148,4 @@ demo = gr.ChatInterface(
 )
 
 # Gradio 데모 실행
-demo.launch(debug=True)
-
+demo.launch(debug=True)