Update app.py

app.py

@@ -26,15 +26,6 @@ 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, {}),
@@ -76,7 +67,6 @@ def load_model(
     return model
 
 
-MAX_NEW_TOKENS = 1500
 EMBEDDER = SentenceTransformer("sentence-transformers/paraphrase-multilingual-mpnet-base-v2")
 MODEL = load_model()
 
@@ -94,11 +84,9 @@ def load_single_document(file_path: str) -> Document:
 
 
 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):
@@ -168,28 +156,28 @@ def bot(
     top_k,
     temp
 ):
+    model = MODEL
     if not history:
         return
 
-    tokens = get_system_tokens(MODEL)[:]
-    tokens.append(LINEBREAK_TOKEN)
+    tokens = get_system_tokens(model)[:]
 
     for user_message, bot_message in history[:-1]:
-        message_tokens = get_message_tokens(model=MODEL, role="user", content=user_message)
+        message_tokens = get_message_tokens(model=model, role="user", content=user_message)
         tokens.extend(message_tokens)
         if bot_message:
-            message_tokens = get_message_tokens(model=MODEL, role="bot", content=bot_message)
+            message_tokens = get_message_tokens(model=model, role="bot", content=bot_message)
             tokens.extend(message_tokens)
 
     last_user_message = history[-1][0]
     if retrieved_docs:
         last_user_message = f"Контекст: {retrieved_docs}\n\nИспользуя контекст, ответь на вопрос: {last_user_message}"
-    message_tokens = get_message_tokens(model=MODEL, role="user", content=last_user_message)
+    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(
+    generator = model.generate(
         tokens,
         top_k=top_k,
         top_p=top_p,
@@ -198,9 +186,9 @@ def bot(
 
     partial_text = ""
     for i, token in enumerate(generator):
-        if token == MODEL.token_eos() or (MAX_NEW_TOKENS is not None and i >= MAX_NEW_TOKENS):
+        if token == model.token_eos():
             break
-        partial_text += MODEL.detokenize([token]).decode("utf-8", "ignore")
+        partial_text += model.detokenize([token]).decode("utf-8", "ignore")
         history[-1][1] = partial_text
         yield history