Update app.py

app.py

@@ -1,37 +1,34 @@
 import os
-import torch
-import faiss
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, Accelerate
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from accelerate import Accelerator  # Accelerate를 별도로 임포트
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
+import faiss
 import gradio as gr
 
-# 환경 변수에서 Hugging Face API 키 로드
-hf_api_key = os.getenv('HF_API_KEY')
+hf_api_key = os.getenv('HF_API_KEY')  # 환경 변수에서 API 키 로드
 
-# 모델 ID 및 토크나이저 설정
 model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key)
+accelerator = Accelerator()  # Accelerator 인스턴스 생성
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     token=hf_api_key,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
     quantization_config=BitsAndBytesConfig(
-        load_in_4bit=True, 
-        bnb_4bit_use_double_quant=True, 
-        bnb_4bit_quant_type="nf4", 
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
         bnb_4bit_compute_dtype=torch.bfloat16
     )
 )
+model = accelerator.prepare(model)  # 모델을 Accelerator에 준비시킴
 
-# 데이터 로딩 및 faiss 인덱스 생성
 ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
 dataset = load_dataset("not-lain/wikipedia", revision="embedded")
 data = dataset["train"]
 data = data.add_faiss_index("embeddings")
 
-# 검색 및 응답 생성 함수
 def search(query: str, k: int = 3):
     embedded_query = ST.encode(query)
     scores, retrieved_examples = data.get_nearest_examples("embeddings", embedded_query, k=k)
@@ -45,8 +42,8 @@ def format_prompt(prompt, retrieved_documents, k):
 
 def generate(formatted_prompt):
     formatted_prompt = formatted_prompt[:2000]  # GPU 메모리 제한을 고려
-    messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]
-    input_ids = tokenizer(messages, return_tensors="pt", padding=True).input_ids.to(model.device)
+    messages = [{"role": "system", "content": "You are an assistant..."}, {"role": "user", "content": formatted_prompt}]
+    input_ids = tokenizer(messages, return_tensors="pt", padding=True).input_ids.to(accelerator.device)
     outputs = model.generate(
         input_ids,
         max_new_tokens=1024,
@@ -65,13 +62,12 @@ def rag_chatbot_interface(prompt: str, k: int = 2):
 
 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,
     inputs=gr.inputs.Textbox(label="Enter your question"),
     outputs=gr.outputs.Textbox(label="Answer"),
     title="Retrieval-Augmented Generation Chatbot",
-    description="This is a chatbot that uses a retrieval-augmented generation approach to provide more accurate answers. It first searches for relevant documents and then generates a response based on the prompt and the retrieved documents."
+    description="This chatbot uses a retrieval-augmented generation approach to provide more accurate answers. It first searches for relevant documents and then generates a response based on the prompt and the retrieved documents."
 )
 
 iface.launch()