Update app.py

app.py

@@ -1,11 +1,13 @@
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from datasets import load_dataset, Dataset
-import faiss  # 필요한 경우 faiss를 임포트합니다.
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import faiss
 import torch
-import os
+import gradio as gr
 
-tokenkey=os.getenv('HF_API_KEY')
+# Hugging Face API 키 환경 변수 설정
+os.environ['HF_API_KEY'] = os.getenv('HF_API_KEY')
 
 # 모델 및 토크나이저 설정
 model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
@@ -31,9 +33,7 @@ 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
-    )
+    scores, retrieved_examples = data.get_nearest_examples("embeddings", embedded_query, k=k)
     return scores, retrieved_examples
 
 def format_prompt(prompt, retrieved_documents, k):
@@ -44,12 +44,12 @@ def format_prompt(prompt, retrieved_documents, k):
 
 def generate(formatted_prompt):
     formatted_prompt = formatted_prompt[:2000]  # GPU 메모리 제한을 고려
-    messages = [{"role": "system", "content": "You are an assistant..."}, {"role": "user", "content": formatted_prompt}]
+    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)
     outputs = model.generate(
         input_ids,
         max_new_tokens=1024,
-        eos_token_id=[tokenizer.eos_token_id],
+        eos_token_id=tokenizer.eos_token_id,
         do_sample=True,
         temperature=0.6,
         top_p=0.9
@@ -57,42 +57,18 @@ def generate(formatted_prompt):
     response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
     return response
 
-def rag_chatbot(prompt: str, k: int = 2):
+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)
 
-rag_chatbot("What is anarchy?", k=2)
-
-
-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)
-
-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."""
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    quantization_config=bnb_config
+# 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."
 )
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
-
-iface = gr.Interface(fn=rag_chatbot_interface, 
-                     inputs="text", 
-                     outputs="text",
-                     input_types=["text"],
-                     output_types=["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."
-                    )
 
-iface.launch()
+iface.launch()