Update app.py

app.py

@@ -10,40 +10,33 @@ from accelerate import Accelerator
 # 환경 변수에서 Hugging Face API 키 로드
 hf_api_key = os.getenv('HF_API_KEY')
 
-# 모델 ID 및 토크나이저 설정
-# 모델 ID
+# 모델 ID 설정
 model_id = "microsoft/phi-2"
 
-# 사용자 정의 코드를 신뢰하고 실행하도록 설정
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    trust_remote_code=True  # 사용자 정의 코드 실행 허용
-)
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key)
-accelerator = Accelerator()
-
-# 양자화 설정 없이 모델 로드 (문제 해결을 위한 임시 조치)
+# 사용자 정의 코드 실행 허용과 함께 모델 및 토크나이저 로드
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_api_key, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     token=hf_api_key,
+    trust_remote_code=True,  # 사용자 정의 코드 실행 허용
     torch_dtype=torch.float32  # 기본 dtype 사용
 )
+
+# Accelerator 설정
+accelerator = Accelerator()
 model = accelerator.prepare(model)
 
-# 데이터 로딩 및 faiss 인덱스 생성
+# 데이터셋 및 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")
 
-# 기타 함수 및 Gradio 인터페이스 구성은 이전과 동일
-
-
-# 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
+    return scores, retrieved_documents
 
 def format_prompt(prompt, retrieved_documents, k):
     PROMPT = f"Question:{prompt}\nContext:"
@@ -52,39 +45,26 @@ def format_prompt(prompt, retrieved_documents, k):
     return PROMPT
 
 def generate(formatted_prompt):
-    # 프롬프트를 문자열로 결합
     prompt_text = f"{SYS_PROMPT} {formatted_prompt}"
-    # 토크나이징
     input_ids = tokenizer(prompt_text, return_tensors="pt", padding=True).input_ids.to(accelerator.device)
-    # 응답 생성
-    outputs = model.generate(
-        input_ids,
-        max_new_tokens=1024,
-        eos_token_id=tokenizer.eos_token_id,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.9
-    )
-    # 응답 텍스트로 디코딩
-    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
-    return response
-
+    outputs = model.generate(input_ids, max_new_tokens=1024, eos_token_id=tokenizer.eos_token_id, do_sample=True, temperature=0.6, top_p=0.9)
+    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."
-
+# 시스템 프롬프트
+SYS_PROMPT = "You are an assistant for answering questions. Provide a conversational answer."
 
+# Gradio 인터페이스
 iface = gr.Interface(
     fn=rag_chatbot_interface,
-    inputs="text",  # 텍스트 입력
-    outputs="text",  # 텍스트 출력
+    inputs="text",
+    outputs="text",
     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 provides more accurate answers by searching relevant documents and generating responses."
 )
 
 iface.launch()