Update app.py

app.py

@@ -12,25 +12,12 @@ from llama_cpp import Llama
 
 
 SYSTEM_PROMPT = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
-SYSTEM_TOKEN = 1587
-USER_TOKEN = 2188
-BOT_TOKEN = 12435
-LINEBREAK_TOKEN = 13
-
-
-ROLE_TOKENS = {
-    "user": USER_TOKEN,
-    "bot": BOT_TOKEN,
-    "system": SYSTEM_TOKEN
-}
 
 
 def get_message_tokens(model, role, content):
-    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
+    content = f"{role}\n{content}\n</s>"
+    content = content.encode("utf-8")
+    return model.tokenize(content, special=True)
 
 
 def get_system_tokens(model):
@@ -51,6 +38,7 @@ model = Llama(
 
 max_new_tokens = 1500
 
+
 def user(message, history):
     new_history = history + [[message, None]]
     return "", new_history
@@ -64,7 +52,6 @@ def bot(
     temp
 ):
     tokens = get_system_tokens(model)[:]
-    tokens.append(LINEBREAK_TOKEN)
 
     for user_message, bot_message in history[:-1]:
         message_tokens = get_message_tokens(model=model, role="user", content=user_message)
@@ -77,7 +64,7 @@ def bot(
     message_tokens = get_message_tokens(model=model, role="user", content=last_user_message)
     tokens.extend(message_tokens)
 
-    role_tokens = [model.token_bos(), BOT_TOKEN, LINEBREAK_TOKEN]
+    role_tokens = model.tokenize("bot\n".encode("utf-8"), special=True)
     tokens.extend(role_tokens)
     generator = model.generate(
         tokens,