Update app.py

app.py

@@ -1,25 +1,20 @@
 import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from sentence_transformers import SentenceTransformer
-from datasets import load_dataset, Dataset
-import faiss
 import torch
+import faiss
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from sentence_transformers import SentenceTransformer
+from datasets import load_dataset
 import gradio as gr
 
-
-# 환경 변수에서 API 키 로드
+# 환경 변수에서 Hugging Face API 키 로드
 hf_api_key = os.getenv('HF_API_KEY')
 
-# 모델 및 토크나이저 로드할 때 API 키 사용
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_api_key)
-model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=hf_api_key)
-
-
-# 모델 및 토크나이저 설정
+# 모델 ID 및 토크나이저 설정
 model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_api_key)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
+    use_auth_token=hf_api_key,
     torch_dtype=torch.bfloat16,
     device_map="auto",
     quantization_config=BitsAndBytesConfig(
@@ -68,6 +63,8 @@ def rag_chatbot_interface(prompt: str, k: int = 2):
     formatted_prompt = format_prompt(prompt, retrieved_documents, k)
     return generate(formatted_prompt)
 
+SYS_PROMPT = "You are an assistant for answering questions. You are given the extracted parts of a long document 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."
+
 # Gradio 인터페이스 설정
 iface = gr.Interface(
     fn=rag_chatbot_interface,