Update app.py

app.py

@@ -11,7 +11,8 @@ import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from pathlib import Path
-
+from pinecone.grpc import PineconeGRPC as Pinecone
+import torch
 from huggingface_hub import CommitScheduler
 
 HF_UPLOAD = os.environ.get("HF_UPLOAD")
@@ -29,6 +30,15 @@ scheduler = CommitScheduler(
     token=HF_UPLOAD
 )
 
+pc = Pinecone(api_key=os.environ.get("PINECONE"))
+index = pc.Index("commonsense")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+retriever_tokenizer = AutoTokenizer.from_pretrained("psyche/dpr-longformer-ko-4096")
+retriever = AutoModel.from_pretrained("psyche/dpr-longformer-ko-4096")
+retriever.eval()
+retriever.to(device)
+
 def save_json(question: str, answer: str) -> None:
     with scheduler.lock:
         with JSON_DATASET_PATH.open("a") as f:
@@ -36,9 +46,9 @@ def save_json(question: str, answer: str) -> None:
             f.write("\n")
 
 
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+MAX_MAX_NEW_TOKENS = 8192
+DEFAULT_MAX_NEW_TOKENS = 4096
+MAX_INPUT_TOKEN_LENGTH = 2048
 
 DESCRIPTION = """\
 # Llama-3 8B Korean QA Chatbot \
@@ -71,8 +81,22 @@ def generate(
         conversation.append({"role": "system", "content": system_prompt})
     for user, assistant in chat_history:
         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-    conversation.append({"role": "user", "content": message})
 
+    
+    retriever_inputs = retriever_tokenizer([message], max_length=1024, truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        embeddings = model(**inputs).pooler_output
+        embeddings = embeddings.cpu().numpy()
+
+    results = index.query(
+        vector=embeddings[0],
+        top_k=1,
+        include_values=False,
+        include_metadata=True
+    )
+    
+    conversation.append({"role": "user", "content": results["matches"][0]["metadata"]+f"\n\n위 문맥을 참고하여 질문 '{message}'에 답하면?"})
+    
     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt", add_generation_prompt=True)
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]