Update app.py

app.py

@@ -1,25 +1,27 @@
 import os
-import torch  # torch를 임포트
-import faiss
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from accelerate import Accelerator
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset
+import faiss
 import gradio as gr
-from accelerate import Accelerator
 
-# 환경 변수에서 Hugging Face API 키 로드
+# Set Hugging Face API key from environment variable
 hf_api_key = os.getenv('HF_API_KEY')
 
-# 모델 ID 및 토크나이저 설정
+# Define model ID
 model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+
+# Initialize tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key)
-accelerator = Accelerator()  # Accelerator 인스턴스 생성
+accelerator = Accelerator()
 
-# 모델 로딩
+# Load the model with custom quantization using BitsAndBytesConfig
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     token=hf_api_key,
-    torch_dtype=torch.bfloat16,  # torch를 사용해 데이터 타입 지정
+    torch_dtype=torch.bfloat16,
     quantization_config=BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_use_double_quant=True,
@@ -27,23 +29,20 @@ model = AutoModelForCausalLM.from_pretrained(
         bnb_4bit_compute_dtype=torch.bfloat16
     )
 )
-model = accelerator.prepare(model)  # 모델을 Accelerator에 준비시킴
+model = accelerator.prepare(model)
 
-# 데이터 로딩 및 faiss 인덱스 생성
+# Load dataset and create FAISS index
 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")
 
-# 검색 및 응답 생성 함수
+# Define functions for search, prompt formatting, and generation
 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):
@@ -51,7 +50,7 @@ def format_prompt(prompt, retrieved_documents, k):
     return PROMPT
 
 def generate(formatted_prompt):
-    formatted_prompt = formatted_prompt[:2000]  # GPU 메모리 제한을 고려
+    formatted_prompt = formatted_prompt[:2000]  # Limit due to GPU memory constraints
     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(
@@ -62,16 +61,17 @@ def generate(formatted_prompt):
         temperature=0.6,
         top_p=0.9
     )
-    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
-    return response
+    return tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
 
 def rag_chatbot_interface(prompt: str, k: int = 2):
     scores, retrieved_documents = search(prompt, k)
     formatted_prompt = format_prompt(prompt, retrieved_documents, k)
     return generate(formatted_prompt)
 
+# Define system prompt for the chatbot
 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."
 
+# Set up Gradio interface
 iface = gr.Interface(
     fn=rag_chatbot_interface,
     inputs=gr.inputs.Textbox(label="Enter your question"),