Update app.py

app.py

@@ -1,54 +1,66 @@
 from sentence_transformers import SentenceTransformer
-from datasets import load_dataset
-import gradio as gr
+from datasets import load_dataset, Dataset
+import faiss  # 필요한 경우 faiss를 임포트합니다.
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
 
-ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
-
-dataset = load_dataset("not-lain/wikipedia",revision = "embedded")
+# 모델 및 토크나이저 설정
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    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", 
+        bnb_4bit_compute_dtype=torch.bfloat16
+    )
+)
 
+# 데이터 로딩 및 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") # column name that has the embeddings of the dataset
+data = data.add_faiss_index("embeddings")
 
-def search(query: str, k: int = 3 ):
-    """a function that embeds a new query and returns the most probable results"""
-    embedded_query = ST.encode(query) # embed new query
-    scores, retrieved_examples = data.get_nearest_examples( # retrieve results
-        "embeddings", embedded_query, # compare our new embedded query with the dataset embeddings
-        k=k # get only top k results
+# 검색 및 응답 생성 함수
+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):
-  """using the retrieved documents we will prompt the model to generate our responses"""
-  PROMPT = f"Question:{prompt}\nContext:"
-  for idx in range(k) :
-    PROMPT+= f"{retrieved_documents['text'][idx]}\n"
-  return PROMPT
+def format_prompt(prompt, retrieved_documents, k):
+    PROMPT = f"Question:{prompt}\nContext:"
+    for idx in range(k):
+        PROMPT += f"{retrieved_documents['text'][idx]}\n"
+    return PROMPT
 
 def generate(formatted_prompt):
-  formatted_prompt = formatted_prompt[:2000] # to avoid GPU OOM
-  messages = [{"role":"system","content":SYS_PROMPT},{"role":"user","content":formatted_prompt}]
-  # tell the model to generate
-  input_ids = tokenizer.apply_chat_template(
-      messages,
-      add_generation_prompt=True,
-      return_tensors="pt"
-  ).to(model.device)
-  outputs = model.generate(
-      input_ids,
-      max_new_tokens=1024,
-      eos_token_id=terminators,
-      do_sample=True,
-      temperature=0.6,
-      top_p=0.9,
-  )
-  response = outputs[0][input_ids.shape[-1]:]
-  return tokenizer.decode(response, skip_special_tokens=True)
+    formatted_prompt = formatted_prompt[:2000]  # GPU 메모리 제한을 고려
+    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(model.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
+
+def rag_chatbot(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(prompt:str,k:int=2):
-  scores , retrieved_documents = search(prompt, k)
-  formatted_prompt = format_prompt(prompt,retrieved_documents,k)
-  return generate(formatted_prompt)
 
 def rag_chatbot_interface(prompt:str,k:int=2):
   scores , retrieved_documents = search(prompt, k)