🌘w🌖

app.py

@@ -22,6 +22,7 @@ device = torch.device("cuda")
 model = model.to(device)
 RAG = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
 TOP_K = 3
+HEADER = "\n# RESOURCES:\n"
 # prepare data
 # since data is too big we will only select the first 3K lines
 
@@ -57,15 +58,14 @@ def prepare_prompt(query, retrieved_examples):
 def talk(message, history):
     retrieved_examples = search(message)
     message, metadata = prepare_prompt(message, retrieved_examples)
-    resources = "\nRESOURCES:\n"
+    resources = HEADER
     for title, url in metadata:
         resources += f"[{title}]({url}),  "
     chat = []
     for item in history:
         chat.append({"role": "user", "content": item[0]})
-        if item[1] is not None:
-            cleaned_past = item[1].split("\nRESOURCES:\n")[0]
-            chat.append({"role": "assistant", "content": cleaned_past})
+        cleaned_past = item[1].split(HEADER)[0]
+        chat.append({"role": "assistant", "content": cleaned_past})
     chat.append({"role": "user", "content": message})
     messages = tokenizer.apply_chat_template(
         chat, tokenize=False, add_generation_prompt=True