app.py

import os
import gradio as gr
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.llms import HuggingFaceHub
from langchain.embeddings import HuggingFaceEmbeddings
import warnings

# Фикс для предупреждений
warnings.filterwarnings("ignore")

# Конфигурация
DOCS_DIR = "lore"
EMBEDDINGS_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
LLM_REPO = "IlyaGusev/saiga_mistral_7b"
HF_TOKEN = os.getenv("HF_TOKEN")  # Добавьте в Secrets Space

# 1. Проверка зависимостей
def check_dependencies():
    try:
        from sentence_transformers import SentenceTransformer
        import torch
        from transformers import pipeline
        print("✔ Все зависимости установлены")
    except ImportError as e:
        raise ImportError(
            f"❌ Не хватает пакетов. Убедитесь, что requirements.txt содержит:\n"
            f"- sentence-transformers\n- torch\n- transformers\n\n"
            f"Ошибка: {str(e)}"
        )

# 2. Загрузка документов
def load_docs():
    docs = []
    if not os.path.exists(DOCS_DIR):
        raise FileNotFoundError(f"Папка {DOCS_DIR} не найдена!")
    
    for file in os.listdir(DOCS_DIR):
        if file.endswith(".txt"):
            try:
                loader = TextLoader(os.path.join(DOCS_DIR, file), encoding="utf-8")
                docs.extend(loader.load())
                print(f"✓ Загружен файл: {file}")
            except Exception as e:
                print(f"⚠ Ошибка в файле {file}: {str(e)}")
    return docs

# 3. Инициализация модели эмбеддингов
def get_embeddings():
    try:
        return HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL)
    except Exception as e:
        raise RuntimeError(f"Ошибка инициализации эмбеддингов: {str(e)}")

# 4. Основная логика
def setup_qa_system():
    check_dependencies()
    
    # Загрузка и обработка документов
    documents = load_docs()
    if not documents:
        raise ValueError("Нет документов для обработки!")
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=300,
        chunk_overlap=30,
        separators=["\n\n", "\n", " ", ""]
    )
    splits = text_splitter.split_documents(documents)
    
    # Создание векторного хранилища
    embeddings = get_embeddings()
    db = FAISS.from_documents(splits, embeddings)
    
    # Инициализация языковой модели
    llm = HuggingFaceHub(
        repo_id=LLM_REPO,
        huggingfacehub_api_token=HF_TOKEN,
        model_kwargs={
            "temperature": 0.2,
            "max_length": 300
        }
    )
    
    return RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=db.as_retriever(search_kwargs={"k": 2}),
        return_source_documents=True
    )

# 5. Функция для интерфейса
def answer_question(question):
    try:
        qa = setup_qa_system()
        result = qa({"query": question})
        answer = result["result"]
        
        # Форматирование ответа
        sources = list({os.path.basename(doc.metadata["source"]) for doc in result["source_documents"]})
        return f"{answer}\n\n(Источники: {', '.join(sources)})"
    except Exception as e:
        return f"⚠ Произошла ошибка: {str(e)}"

# Интерфейс
with gr.Blocks(title="📚 Лор-бот") as app:
    gr.Markdown("## 🧛 Справочник по сверхъестественному")
    with gr.Row():
        question = gr.Textbox(label="Ваш вопрос", placeholder="Какие слабости у вампиров?")
        submit = gr.Button("Спросить")
    answer = gr.Textbox(label="Ответ", interactive=False)
    
    submit.click(answer_question, inputs=question, outputs=answer)

app.launch(server_name="0.0.0.0", server_port=7860)