app.py

import gradio as gr
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_community.llms import HuggingFaceHub
from langchain.prompts import ChatPromptTemplate
from dotenv import load_dotenv
import os
import
 shutil

# Load environment variables
load_dotenv()

CHROMA_PATH = "/tmp/chroma"
DATA_PATH = ""  # Укажите путь к вашим данным
PROMPT_TEMPLATE = """
Ответь на вопрос, используя только следующий контекст:
{context}
---
Ответь на вопрос на основе приведенного контекста: {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 generate_data_store():
    documents = load_documents()
    if documents:
        chunks = split_text(documents)
        save_to_chroma(chunks)

def load_documents():
    file_path = os.path.join(DATA_PATH, "pl250320252.md")
    if not os.path.exists(file_path):
        print(f"Ошибка: Файл {file_path} не найден.")
        return []
    loader = UnstructuredMarkdownLoader(file_path)
    documents = loader.load()
    return documents

def split_text(documents: list[Document]):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=900,
        chunk_overlap=300,
        length_function=len,
        add_start_index=True,
    )
    chunks = text_splitter.split_documents(documents)
    print(f"Разделено {len(documents)} документов на {len(chunks)} частей.")
    return chunks


def save_to_chroma(chunks: list[Document]):
    # Очищаем базу данных перед сохранением
    if os.path.exists(CHROMA_PATH):
        shutil.rmtree(CHROMA_PATH)

    # Инициализация эмбеддингов
    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}
    )

    # Создание Chroma DB
    db = Chroma.from_documents(
        chunks,
        embeddings,
        persist_directory=CHROMA_PATH
    )
    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 'Нет источников'}"

# Создаем папки при необходимости
os.makedirs("/tmp/model_cache", exist_ok=True)
os.makedirs("/tmp/chroma", exist_ok=True)

# Создание и запуск интерфейса Gradio
interface = gr.Interface(
    fn=chat_interface,
    inputs=gr.Textbox(lines=2, placeholder="Введите ваш вопрос здесь..."),
    outputs="text",
    title="Чат с документами",
    description="Задайте вопрос, и я отвечу на основе загруженных документов."
)

if __name__ == "__main__":
    interface.launch()