Update app.py

app.py

@@ -26,6 +26,15 @@ from llama_cpp import Llama
 
 
 SYSTEM_PROMPT = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
+SYSTEM_TOKEN = 1788
+USER_TOKEN = 1404
+BOT_TOKEN = 9225
+LINEBREAK_TOKEN = 13
+ROLE_TOKENS = {
+    "user": USER_TOKEN,
+    "bot": BOT_TOKEN,
+    "system": SYSTEM_TOKEN
+}
 
 LOADER_MAPPING = {
     ".csv": (CSVLoader, {}),
@@ -85,9 +94,11 @@ def load_single_document(file_path: str) -> Document:
 
 
 def get_message_tokens(model, role, content):
-    content = f"{role}\n{content}\n</s>"
-    content = content.encode("utf-8")
-    return model.tokenize(content, special=True)
+    message_tokens = model.tokenize(content.encode("utf-8"))
+    message_tokens.insert(1, ROLE_TOKENS[role])
+    message_tokens.insert(2, LINEBREAK_TOKEN)
+    message_tokens.append(model.token_eos())
+    return message_tokens
 
 
 def get_system_tokens(model):