Update app.py

app.py

@@ -52,9 +52,11 @@ def format_prompt(prompt, retrieved_documents, k):
     return PROMPT
 
 def generate(formatted_prompt):
-    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)
+    # 프롬프트를 문자열로 결합
+    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,
@@ -63,7 +65,10 @@ def generate(formatted_prompt):
         temperature=0.6,
         top_p=0.9
     )
-    return tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
+    # 응답 텍스트로 디코딩
+    response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
+    return response
+
 
 def rag_chatbot_interface(prompt: str, k: int = 2):
     scores, retrieved_documents = search(prompt, k)