Update app.py

app.py

@@ -1,117 +1,168 @@
-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
+from sentence_transformers import SentenceTransformer
+import chromadb
+from chromadb.utils import embedding_functions
+import os
+from langdetect import detect
 
-# Фикс для предупреждений
-warnings.filterwarnings("ignore")
+# Проверяем наличие текстовых файлов и читаем их
+def load_text_files():
+    files = {
+        "vampires": "vampires.txt",
+        "werewolves": "werewolves.txt",
+        "humans": "humans.txt"
+    }
+    
+    loaded_data = {}
+    for key, filename in files.items():
+        try:
+            with open(filename, 'r', encoding='utf-8') as file:
+                loaded_data[key] = file.read()
+        except FileNotFoundError:
+            print(f"Файл {filename} не найден")
+            loaded_data[key] = ""
+    
+    return loaded_data
 
-# Конфигурация
-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
+# Инициализация модели для эмбеддингов
+def initialize_embedding_model():
+    return embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+    )
 
-# 1. Проверка зависимостей
-def check_dependencies():
+# Создание базы знаний
+def create_knowledge_base(text_data, embed_fn):
+    client = chromadb.Client()
+    
     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)}"
+        collection = client.get_collection(name="knowledge_base")
+    except:
+        collection = client.create_collection(
+            name="knowledge_base",
+            embedding_function=embed_fn
         )
+    
+    # Добавляем документы в коллекцию
+    documents = []
+    metadatas = []
+    ids = []
+    
+    for category, text in text_data.items():
+        if text:  # только если текст не пустой
+            # Разбиваем текст на предложения или абзацы
+            paragraphs = [p for p in text.split('\n') if p.strip()]
+            
+            for i, paragraph in enumerate(paragraphs):
+                documents.append(paragraph)
+                metadatas.append({"category": category})
+                ids.append(f"{category}_{i}")
+    
+    if documents:
+        collection.add(
+            documents=documents,
+            metadatas=metadatas,
+            ids=ids
+        )
+    
+    return collection
 
-# 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", " ", ""]
+# Инициализация модели для ответов
+def initialize_llm_model():
+    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+    
+    model_name = "IlyaGusev/saiga_mistral_7b"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device="cpu"
     )
-    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 pipe
+
+# Поиск релевантной информации
+def find_relevant_info(question, collection, embed_fn, n_results=3):
+    results = collection.query(
+        query_texts=[question],
+        n_results=n_results
     )
     
-    return RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=db.as_retriever(search_kwargs={"k": 2}),
-        return_source_documents=True
+    context = "\n\n".join(results['documents'][0])
+    return context
+
+# Генерация ответа
+def generate_response(question, context, llm_pipe):
+    system_prompt = """Ты - помощник, который отвечает на вопросы пользователя, используя предоставленную информацию. 
+    Отвечай только на основе предоставленного контекста. Если ответа нет в контексте, скажи, что не знаешь.
+    Отвечай на русском языке."""
+    
+    prompt = f"""<s>{system_prompt}
+    Контекст: {context}
+    Вопрос: {question}
+    Ответ:"""
+    
+    output = llm_pipe(
+        prompt,
+        max_new_tokens=512,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        eos_token_id=2
     )
+    
+    return output[0]["generated_text"][len(prompt):].strip()
 
-# 5. Функция для интерфейса
-def answer_question(question):
+# Основная функция для обработки запросов
+def answer_question(question, history):
+    # Определяем язык вопроса
     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)}"
+        lang = detect(question)
+        if lang != 'ru':
+            return "Пожалуйста, задавайте вопросы на русском языке."
+    except:
+        pass
+    
+    # Загружаем данные (если еще не загружены)
+    if not hasattr(answer_question, 'text_data'):
+        answer_question.text_data = load_text_files()
+    
+    # Инициализируем модели (если еще не инициализированы)
+    if not hasattr(answer_question, 'embed_fn'):
+        answer_question.embed_fn = initialize_embedding_model()
+    
+    if not hasattr(answer_question, 'collection'):
+        answer_question.collection = create_knowledge_base(answer_question.text_data, answer_question.embed_fn)
+    
+    if not hasattr(answer_question, 'llm_pipe'):
+        answer_question.llm_pipe = initialize_llm_model()
+    
+    # Находим релевантный контекст
+    context = find_relevant_info(question, answer_question.collection, answer_question.embed_fn)
+    
+    # Генерируем ответ
+    response = generate_response(question, context, answer_question.llm_pipe)
+    
+    return response
 
-# Интерфейс
-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)
+# Создаем интерфейс Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("## Чат-бот с доступом к текстовым файлам")
+    gr.Markdown("Задавайте вопросы о вампирах, оборотнях или людях на русском языке")
+    
+    chatbot = gr.Chatbot(label="Диалог")
+    msg = gr.Textbox(label="Ваш вопрос")
+    clear = gr.Button("Очистить")
+    
+    def respond(message, chat_history):
+        bot_message = answer_question(message, chat_history)
+        chat_history.append((message, bot_message))
+        return "", chat_history
+    
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
 
-app.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()