Update app.py

app.py

@@ -18,10 +18,7 @@ from langchain.document_loaders import (
     UnstructuredWordDocumentLoader,
 )
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.docstore.document import Document
-from chromadb.config import Settings
 from llama_cpp import Llama
 
 
@@ -66,9 +63,9 @@ def load_model(
     print("Model loaded!")
     return model
 
+
 MAX_NEW_TOKENS = 1500
-EMBEDDER_NAME = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
-EMBEDDER = HuggingFaceEmbeddings(model_name=EMBEDDER_NAME)
+EMBEDDER = SentenceTransformer("sentence-transformers/paraphrase-multilingual-mpnet-base-v2")
 MODEL = load_model()
 
 
@@ -121,15 +118,12 @@ def build_index(file_paths, db, chunk_size, chunk_overlap, file_warning):
             continue
         fixed_documents.append(doc)
     print("Documents after processing:", len(fixed_documents))
-    
-    db = Chroma.from_documents(
-        fixed_documents,
-        EMBEDDER,
-        client_settings=Settings(
-            anonymized_telemetry=False
-        )
-    )
+
+    texts = [doc.page_content for doc in fixed_documents]
+    embeddings = EMBEDDER.encode(texts, convert_to_tensor=True)
+    db = {"docs": texts, "embeddings": embeddings}
     print("Embeddings calculated!")
+    
     file_warning = f"Загружено {len(fixed_documents)} фрагментов! Можно задавать вопросы."
     return db, file_warning
 
@@ -138,9 +132,11 @@ def retrieve(history, db, retrieved_docs, k_documents):
     retrieved_docs = ""
     if db:
         last_user_message = history[-1][0]
-        retriever = db.as_retriever(search_kwargs={"k": k_documents})
-        docs = retriever.get_relevant_documents(last_user_message)
-        retrieved_docs = "\n\n".join([doc.page_content for doc in docs])
+        query_embedding = EMBEDDER.encode(last_user_message, convert_to_tensor=True)
+        scores = cos_sim(query_embedding, db["embeddings"])[0]
+        top_k_idx = torch.topk(scores, k=k_documents)[1]
+        top_k_documents = [db["docs"][idx] for idx in top_k_idx]
+        retrieved_docs = "\n\n".join(top_k_documents)
     return retrieved_docs