Interact fix

interact.py

@@ -1,15 +1,23 @@
 import fire
 from llama_cpp import Llama
 
-MESSAGE_TEMPLATE = "{role}\n{content}"
 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
+}
+
 
 def get_message_tokens(model, role, content):
-    message_text = MESSAGE_TEMPLATE.format(role=role, content=content)
-    message_tokens = model.tokenize(message_text.encode("utf-8"))
+    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