Update app.py

app.py

@@ -1,3 +1,4 @@
+import gradio as gr
 from langchain_community.document_loaders import UnstructuredMarkdownLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_core.documents import Document
@@ -12,7 +13,7 @@ import shutil
 # Load environment variables
 load_dotenv()
 
-CHROMA_PATH = "chroma"
+CHROMA_PATH = "/tmp/chroma"
 DATA_PATH = ""  # Укажите путь к вашим данным
 PROMPT_TEMPLATE = """
 Ответь на вопрос, используя только следующий контекст:
@@ -21,10 +22,26 @@ PROMPT_TEMPLATE = """
 Ответь на вопрос на основе приведенного контекста: {question}
 """
 
+# Инициализация базы данных Chroma один раз при запуске
+def initialize_chroma():
+    if not os.path.exists(CHROMA_PATH):
+        generate_data_store()
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-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
+
 def main():
     # Создаем папки при необходимости
-    os.makedirs("model_cache", exist_ok=True)
-    os.makedirs("chroma", exist_ok=True)
+    os.makedirs("/tmp/model_cache", exist_ok=True)
+    os.makedirs("/tmp/chroma", exist_ok=True)
 
     # Генерируем или загружаем базу данных Chroma один раз
     if not os.path.exists(CHROMA_PATH):
@@ -52,33 +69,6 @@ def main():
             break
         process_query(query_text, db)
 
-def process_query(query_text: str, db):
-    # Поиск по схожести
-    results = db.similarity_search_with_relevance_scores(query_text, k=3)
-    if not results or results[0][1] < 0.7:
-        print("Не найдено подходящих результатов.")
-        return
-
-    # Формирование контекста
-    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]
-    print(f"Ответ: {response_text}")
-    print(f"Источники: {sources}")
 
 def generate_data_store():
     documents = load_documents()
@@ -114,7 +104,7 @@ def save_to_chroma(chunks: list[Document]):
     # Инициализация эмбеддингов
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-        cache_folder="model_cache",
+        cache_folder="/tmp/model_cache",
         model_kwargs={'device': 'cpu'},
         encode_kwargs={'normalize_embeddings': True}
     )
@@ -127,5 +117,36 @@ def save_to_chroma(chunks: list[Document]):
     )
     print(f"Сохранено {len(chunks)} частей в {CHROMA_PATH}.")
 
+# Обработка запроса пользователя
+def process_query(query_text: str, db):
+    results = db.similarity_search_with_relevance_scores(query_text, k=3)
+    if not results or results[0][1] < 0.7:
+        return "Не найдено подходящих результатов.", []
+    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
+
+# Функция для интерфейса Gradio
+def chat_interface(query_text):
+    db = initialize_chroma()
+    response, sources = process_query(query_text, db)
+    return f"Ответ: {response}\n\nИсточники: {', '.join(sources) if sources else 'Нет источников'}"
+
+# Создание и запуск интерфейса Gradio
+interface = gr.Interface(
+    fn=chat_interface,
+    inputs=gr.Textbox(lines=2, placeholder="Введите ваш вопрос здесь..."),
+    outputs="text",
+    title="Чат с документами",
+    description="Задайте вопрос, и я отвечу на основе загруженных документов."
+)
+
 if __name__ == "__main__":
     main()