Update app.py

app.py

@@ -1,20 +1,25 @@
 import os
+import torch  # torch를 임포트
+import faiss
 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
+from accelerate import Accelerator
 
-hf_api_key = os.getenv('HF_API_KEY')  # 환경 변수에서 API 키 로드
+# 환경 변수에서 Hugging Face API 키 로드
+hf_api_key = os.getenv('HF_API_KEY')
 
+# 모델 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,
+    torch_dtype=torch.bfloat16,  # torch를 사용해 데이터 타입 지정
     quantization_config=BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_use_double_quant=True,
@@ -24,16 +29,21 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 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)
     return scores, retrieved_examples
 
+# 나머지 코드는 이전과 동일하게 유지
+
+
 def format_prompt(prompt, retrieved_documents, k):
     PROMPT = f"Question:{prompt}\nContext:"
     for idx in range(k):