Update app.py

app.py

@@ -26,16 +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, {}),
@@ -52,37 +42,42 @@ LOADER_MAPPING = {
     ".txt": (TextLoader, {"encoding": "utf8"}),
 }
 
-directory = "."
-model_url = "https://huggingface.co/IlyaGusev/saiga2_13b_gguf/resolve/main/model-q4_K.gguf"
-repo_name = "IlyaGusev/saiga2_13b_gguf"
-model_name = "model-q4_K.gguf"
-final_model_path = os.path.join(directory, model_name)
-embedder_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-
-print("Downloading all files...")
-rm_files = [os.path.join(directory, f) for f in os.listdir(directory)]
-for f in rm_files:
-   if os.path.isfile(f):
-       os.remove(f)
-   else:
-       shutil.rmtree(f)
-
-if not os.path.exists(final_model_path):
-    with open(final_model_path, "wb") as f:
-        http_get(model_url, f)
-os.chmod(final_model_path, 0o777)
-print("Files downloaded!")
-
-model = Llama(
-    model_path=final_model_path,
-    n_ctx=2000,
-    n_parts=1,
-)
 
-print("Model loaded!")
+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"
+):
+    final_model_path = os.path.join(directory, model_name)
+    
+    print("Downloading all files...")
+    rm_files = [os.path.join(directory, f) for f in os.listdir(directory)]
+    for f in rm_files:
+       if os.path.isfile(f):
+           os.remove(f)
+       else:
+           shutil.rmtree(f)
+    
+    if not os.path.exists(final_model_path):
+        with open(final_model_path, "wb") as f:
+            http_get(model_url, f)
+    os.chmod(final_model_path, 0o777)
+    print("Files downloaded!")
+    
+    model = Llama(
+        model_path=final_model_path,
+        n_ctx=2000,
+        n_parts=1,
+    )
+    
+    print("Model loaded!")
+    return model
+
+MAX_NEW_TOKENS = 1500
+EMBEDDER_NAME = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
+EMBEDDER = HuggingFaceEmbeddings(model_name=EMBEDDER_NAME)
+MODEL = load_model()
 
-max_new_tokens = 1500
-embeddings = HuggingFaceEmbeddings(model_name=embedder_name)
 
 def get_uuid():
     return str(uuid4())
@@ -97,11 +92,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):
@@ -136,7 +129,7 @@ def build_index(file_paths, db, chunk_size, chunk_overlap, file_warning):
 
     db = Chroma.from_documents(
         fixed_documents,
-        embeddings,
+        EMBEDDER,
         client_settings=Settings(
             anonymized_telemetry=False
         )
@@ -151,7 +144,7 @@ def user(message, history, system_prompt):
 
 
 def retrieve(history, db, retrieved_docs, k_documents):
-    context = ""
+    retrieved_docs = ""
     if db:
         last_user_message = history[-1][0]
         retriever = db.as_retriever(search_kwargs={"k": k_documents})
@@ -172,25 +165,25 @@ def bot(
     if not history:
         return
 
-    tokens = get_system_tokens(model)[:]
+    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)
+        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.token_bos(), BOT_TOKEN, LINEBREAK_TOKEN]
     tokens.extend(role_tokens)
-    generator = model.generate(
+    generator = MODEL.generate(
         tokens,
         top_k=top_k,
         top_p=top_p,
@@ -199,9 +192,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() or (MAX_NEW_TOKENS is not None and i >= MAX_NEW_TOKENS):
             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