Update app.py

app.py

@@ -9,21 +9,10 @@ from llama_cpp import Llama
 SYSTEM_PROMPT = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
 
 
-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)
-
-
-def get_system_tokens(model):
-    system_message = {"role": "system", "content": SYSTEM_PROMPT}
-    return get_message_tokens(model, **system_message)
-
-
 def load_model(
     directory: str = ".",
-    model_name: str = "model-q4_K.gguf",
-    model_url: str = "https://huggingface.co/IlyaGusev/saiga2_13b_gguf/resolve/main/model-q4_K.gguf"
+    model_name: str = "saiga_nemo_12b.Q4_K_M.gguf",
+    model_url: str = "https://huggingface.co/IlyaGusev/saiga_nemo_12b_gguf/resolve/main/saiga_nemo_12b.Q4_K_M.gguf"
 ):
     final_model_path = os.path.join(directory, model_name)
     
@@ -36,7 +25,7 @@ def load_model(
     
     model = Llama(
         model_path=final_model_path,
-        n_ctx=1024
+        n_ctx=8192
     )
     
     print("Model loaded!")
@@ -59,35 +48,28 @@ def bot(
     temp
 ):
     model = MODEL
-    tokens = get_system_tokens(model)[:]
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
 
     for user_message, bot_message in history[:-1]:
-        message_tokens = get_message_tokens(model=model, role="user", content=user_message)
-        tokens.extend(message_tokens)
+        messages.append({"role": "user", "content": user_message})
         if bot_message:
-            message_tokens = get_message_tokens(model=model, role="bot", content=bot_message)
-            tokens.extend(message_tokens)
+            messages.append({"role": "assistant", "content": bot_message})
 
     last_user_message = history[-1][0]
-    message_tokens = get_message_tokens(model=model, role="user", content=last_user_message)
-    tokens.extend(message_tokens)
-
-    role_tokens = model.tokenize("bot\n".encode("utf-8"), special=True)
-    tokens.extend(role_tokens)
-    generator = model.generate(
-        tokens,
+    messages.append({"role": "user", "content": last_user_message})
+    partial_text = ""
+    for part in model.create_chat_completion(
+        messages,
+        temperature=temp,
         top_k=top_k,
         top_p=top_p,
-        temp=temp
-    )
-
-    partial_text = ""
-    for i, token in enumerate(generator):
-        if token == model.token_eos():
-            break
-        partial_text += model.detokenize([token]).decode("utf-8", "ignore")
-        history[-1][1] = partial_text
-        yield history
+        stream=True,
+    ):
+        delta = part["choices"][0]["delta"]
+        if "content" in delta:
+            partial_text += delta["content"]
+            history[-1][1] = partial_text
+            yield history
 
 
 with gr.Blocks(