Update app.py

app.py

@@ -1,7 +1,7 @@
 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
@@ -11,7 +11,6 @@ import fitz  # PyMuPDF
 # 환경 변수에서 Hugging Face 토큰 가져오기
 token = os.environ.get("HF_TOKEN")
 
-
 # 임베딩 모델 로드
 ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
 
@@ -43,17 +42,13 @@ data = dataset["train"]
 data = data.map(lambda x: {"question_embedding": ST.encode(x["question"])}, batched=True)
 data.add_faiss_index(column="question_embedding")
 
-# LLaMA 모델 설정
+# LLaMA 모델 설정 (양자화 없이)
 model_id = "google/gemma-2-2b-it"
-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 = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.bfloat16,
+    torch_dtype=torch.bfloat16,  # 양자화 없이 bfloat16 사용
     device_map="auto",
-    quantization_config=bnb_config,
     token=token
 )
 
@@ -62,7 +57,6 @@ You are given the extracted parts of legal documents and a question. Provide a c
 If you don't know the answer, just say "I do not know." Don't make up an answer.
 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.
-
 Example 1:
 User: I need help with a contract dispute.
 Assistant: Hello! I'm your friendly GPT legal advisor bot (v0.1.0), and I'm here to help you with your contract dispute by asking you a series of questions. You can ask for help, more details, or a summary at any time. Let's get started! What is the nature of the contract in question?
@@ -72,7 +66,6 @@ User: California
 Assistant: When did you enter into the lease agreement?
 User: January 1st, 2022
 ...
-
 Example 2:
 User: I need help with a copyright issue.
 Assistant: Hi there! I'm your legal advisor bot, and I'll be assisting you with your copyright issue by asking some questions. You can request help, more details, or a summary at any time. Let's begin! What is the copyrighted material you are concerned about?
@@ -184,7 +177,8 @@ Damages: Present evidence of the financial harm you suffered due to the infringe
 Injunction: Request a court order to stop the defendant from continuing to infringe on your intellectual property rights.
 ...
 
-Begin by introducing yourself, next tell them they can ask for help or more details or a summary at any time, and start by asking what they need help with."""
+Begin by introducing yourself, next tell them they can ask for help or more details or a summary at any time, and start by asking what they need help with.
+you must answer korean."""
 
 # 법률 문서 검색 함수
 def search_law(query, k=5):
@@ -203,7 +197,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"  # Assuming doc[0] contains the relevant text
+        PROMPT += f"{doc[0]}\n"
     PROMPT += "\nLegal QA:\n"
     for doc in qa_docs:
         PROMPT += f"{doc}\n"
@@ -271,4 +265,4 @@ demo = gr.ChatInterface(
 )
 
 # Gradio 데모 실행
-demo.launch(debug=True)
+demo.launch(debug=True)