Update app.py

app.py

@@ -3,19 +3,16 @@ from langchain_community.document_loaders import UnstructuredMarkdownLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_core.documents import Document
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_chroma import Chroma
-from langchain.vectorstores import FAISS
+from langchain_community.vectorstores import FAISS  # Обновленный импорт
 from langchain_community.llms import HuggingFaceHub
 from langchain.prompts import ChatPromptTemplate
 from dotenv import load_dotenv
 import os
-import shutil
 
 # Загрузка переменных окружения
 load_dotenv()
 
-CHROMA_PATH = "/tmp/chroma"
-DATA_PATH = ""  # Укажите путь к вашим данным, например "data", если файл не в корне
+DATA_PATH = ""  # Укажите путь к вашему файлу
 PROMPT_TEMPLATE = """
 Ответь на вопрос, используя только следующий контекст:
 {context}
@@ -26,47 +23,29 @@ PROMPT_TEMPLATE = """
 # Глобальная переменная для статуса
 status_message = "Инициализация..."
 
-def initialize_chroma():
+def initialize_vectorstore():
     global status_message
-    status_message = "Инициализация векторной базы..."
-    else:
-        status_message = "База данных Chroma уже существует. Пересоздаем базу данных..."
-        shutil.rmtree(CHROMA_PATH)
+    try:
+        status_message = "Загрузка и обработка документов..."
+        documents = load_documents()
+        chunks = split_text(documents)
         
-    generate_data_store()
-    status_message
- = "База данных Chroma создана и подготовлена."
-    
-    embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
-        cache_folder="/tmp/model_cache",
-        model_kwargs={'device': 'cpu'},
-        encode_kwargs={'normalize_embeddings': True}
-    )
-    db = Chroma(
-        persist_directory=CHROMA_PATH,
-        embedding_function=embeddings
-    )
-    return db
+        status_message = "Создание векторной базы..."
+        vectorstore = save_to_faiss(chunks)
+        
+        status_message = "База данных готова к использованию."
+        return vectorstore
+        
+    except Exception as e:
+        status_message = f"Ошибка инициализации: {str(e)}"
+        raise
 
-def generate_data_store():
-    documents = load_documents()
-    if not documents:
-        raise Exception("Документы не загружены!")
-    chunks = split_text(documents)
-    status_message = f"Первые 50 символов первого чанка: {chunks[0].page_content[:50]}"
-    save_to_chroma(chunks)
-    
-    
 def load_documents():
     file_path = os.path.join(DATA_PATH, "pl250320252.md")
     if not os.path.exists(file_path):
-        global status_message
-        status_message = f"Ошибка: Файл {file_path} не найден."
-        return []
+        raise FileNotFoundError(f"Файл {file_path} не найден")
     loader = UnstructuredMarkdownLoader(file_path)
-    documents = loader.load()
-    return documents
+    return loader.load()
 
 def split_text(documents: list[Document]):
     text_splitter = RecursiveCharacterTextSplitter(
@@ -75,59 +54,58 @@ def split_text(documents: list[Document]):
         length_function=len,
         add_start_index=True,
     )
-    chunks = text_splitter.split_documents(documents)
-    global status_message
-    status_message = f"\nРазделено {len(documents)} документов на {len(chunks)} частей."
-    return chunks
+    return text_splitter.split_documents(documents)
 
-def save_to_chroma(chunks: list[Document]):
-    if os.path.exists(CHROMA_PATH):
-        shutil.rmtree(CHROMA_PATH)
+def save_to_faiss(chunks: list[Document]):
     embeddings = HuggingFaceEmbeddings(
-        model_name="sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
-        model_kwargs={'device': 'cpu', 'trust_remote_code': True},
-        encode_kwargs={'normalize_embeddings': False}  # Экспериментируйте с этим
-    )
-    Chroma.from_documents(
-        chunks,
-        embeddings,
-        persist_directory=CHROMA_PATH
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        model_kwargs={'device': 'cpu'},
+        encode_kwargs={'normalize_embeddings': True}
     )
+    return FAISS.from_documents(chunks, embeddings)
 
-def process_query(query_text: str, db):
-    results = db.similarity_search_with_relevance_scores(query_text, k=3)
-    global status_message
-    status_message += f"\nНайдено {len(results)} результатов с релевантностью: {[round(score, 2) for _, score in results]}"
-    
-    if not results:
-        return "Не найдено подходящих результатов.", []
+def process_query(query_text: str, vectorstore):
+    if vectorstore is None:
+        return "База данных не инициализирована", []
     
-    if results[0][1] < 0.5:
-        response = "Нет высококачественных совпадений. Возможные варианты:\n"
-        response += "\n".join([f"- {doc.page_content[:50]}..." for doc, _ in results])
-        return response, [doc.metadata.get("source") for doc, _ in results]
-    
-    context_text = "\n\n---\n\n".join([doc.page_content for doc, _ in results])
-    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
-    prompt = prompt_template.format(context=context_text, question=query_text)
-    model = HuggingFaceHub(
-        repo_id="google/flan-t5-small",
-        model_kwargs={"temperature": 0.5, "max_length": 512}
-    )
-    response_text = model.predict(prompt)
-    sources = [doc.metadata.get("source", None) for doc, _ in results]
-    return response_text, sources
+    try:
+        results = vectorstore.similarity_search_with_relevance_scores(query_text, k=3)
+        global status_message
+        status_message += f"\nНайдено {len(results)} результатов"
+        
+        if not results:
+            return "Не найдено результатов.", []
+            
+        context_text = "\n\n---\n\n".join([
+            f"Релевантность: {score:.2f}\n{doc.page_content}" 
+            for doc, score in results
+        ])
+        
+        prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
+        prompt = prompt_template.format(context=context_text, question=query_text)
+        
+        model = HuggingFaceHub(
+            repo_id="bert-base-multilingual-cased",
+            task="text2text-generation",  # Указываем задачу
+            model_kwargs={"temperature": 0.5, "max_length": 512}
+        )
+        response_text = model.predict(prompt)
+        
+        sources = list(set([doc.metadata.get("source", "") for doc, _ in results]))
+        return response_text, sources
+        
+    except Exception as e:
+        return f"Ошибка обработки запроса: {str(e)}", []
 
 def chat_interface(query_text):
     global status_message
-    db = initialize_chroma()
-    response, sources = process_query(query_text, db)
-    full_response = f"{status_message}\n\nОтвет: {response}\n\nИсточники: {', '.join(sources) if sources else 'Нет источников'}"
-    return full_response
-
-# Создание папок
-os.makedirs("/tmp/model_cache", exist_ok=True)
-os.makedirs("/tmp/chroma", exist_ok=True)
+    try:
+        vectorstore = initialize_vectorstore()
+        response, sources = process_query(query_text, vectorstore)
+        full_response = f"{status_message}\n\nОтвет: {response}\n\nИсточники: {', '.join(sources) if sources else 'Нет источников'}"
+        return full_response
+    except Exception as e:
+        return f"Критическая ошибка: {str(e)}"
 
 # Интерфейс Gradio
 interface = gr.Interface(