Update app.py

app.py

@@ -1,90 +1,97 @@
 import os
-import re
-import torch  # Добавлен импорт torch
+import gradio as gr
 from langchain_community.document_loaders import TextLoader
 from langchain.text_splitter import CharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_core.prompts import PromptTemplate
 from langchain.chains import RetrievalQA
-from transformers import pipeline
-import gradio as gr
+from langchain_community.llms import HuggingFacePipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# 1. Загрузка всех файлов из папки lore/
-def load_all_lore_files():
+# Конфигурация
+DOCS_DIR = "lore"
+MODEL_NAME = "IlyaGusev/saiga_mistral_7b"  # Оптимальная модель для русского
+EMBEDDINGS_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+
+# 1. Загрузка документов
+def load_documents():
     docs = []
-    for filename in os.listdir("lore"):
+    for filename in os.listdir(DOCS_DIR):
         if filename.endswith(".txt"):
-            loader = TextLoader(os.path.join("lore", filename), encoding="utf-8")
+            loader = TextLoader(os.path.join(DOCS_DIR, filename), encoding="utf-8")
             docs.extend(loader.load())
     return docs
 
-# 2. Очистка от спецсимволов
-def clean_text(text):
-    return re.sub(r"\[=.*?\/?]", "", text)
+# 2. Подготовка базы знаний
+def prepare_knowledge_base():
+    documents = load_documents()
+    
+    # Разбиваем текст на чанки
+    text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+    splits = text_splitter.split_documents(documents)
+    
+    # Создаем векторное хранилище
+    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDINGS_MODEL)
+    vectorstore = FAISS.from_documents(splits, embeddings)
+    
+    return vectorstore
 
-# 3. Настройка эмбеддингов
-def create_embeddings():
-    return HuggingFaceEmbeddings(
-        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+# 3. Инициализация языковой модели
+def load_llm():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        load_in_4bit=True  # Экономия памяти
     )
-
-# 4. Создание векторной базы
-def create_vectorstore(docs, embeddings):
-    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-    split_docs = text_splitter.split_documents(docs)
-    for doc in split_docs:
-        doc.page_content = clean_text(doc.page_content)
-    return FAISS.from_documents(split_docs, embeddings)
-
-# 5. Загрузка модели ответа (с проверкой доступности GPU)
-def create_llm_pipeline():
-    return pipeline(
+    
+    pipe = pipeline(
         "text-generation",
-        model="IlyaGusev/saiga2_7b_lora",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device="cuda:0" if torch.cuda.is_available() else "cpu"
-    )
-
-# 6. Объединение в цепочку
-def build_chain():
-    docs = load_all_lore_files()
-    embeddings = create_embeddings()
-    vectorstore = create_vectorstore(docs, embeddings)
-
-    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-
-    prompt = PromptTemplate(
-        template="""
-Ты — помощник, который отвечает на вопросы по вымышленному лору. Отвечай кратко, точно и на русском языке.
-Если в лоре нет нужной информации, честно скажи, что не знаешь.
-
-Контекст:
-{context}
-
-Вопрос:
-{question}
-
-Ответ:
-""",
-        input_variables=["context", "question"]
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=200,
+        temperature=0.3
     )
+    
+    return HuggingFacePipeline(pipeline=pipe)
 
+# 4. Создание цепочки для вопросов-ответов
+def create_qa_chain():
+    vectorstore = prepare_knowledge_base()
+    llm = load_llm()
+    
     return RetrievalQA.from_chain_type(
-        llm=create_llm_pipeline(),
-        retriever=retriever,
-        chain_type_kwargs={"prompt": prompt}
+        llm=llm,
+        chain_type="stuff",
+        retriever=vectorstore.as_retriever(search_kwargs={"k": 2}),
+        return_source_documents=True
     )
 
-# 7. Интерфейс
-qa_chain = build_chain()
+# 5. Функция для ответов
+def get_answer(question):
+    qa_chain = create_qa_chain()
+    result = qa_chain({"query": question})
+    
+    # Форматируем ответ
+    answer = result["result"]
+    sources = list(set(doc.metadata["source"] for doc in result["source_documents"]))
+    
+    return f"{answer}\n\nИсточники: {', '.join(sources)}"
 
-def ask_question(question):
-    return qa_chain.run(question)
+# 6. Интерфейс Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧛 Лор-бот: справочник по сверхъестественному")
+    
+    with gr.Row():
+        question = gr.Textbox(label="Ваш вопрос", placeholder="Какие слабости у вампиров?")
+        submit_btn = gr.Button("Спросить")
+    
+    answer = gr.Textbox(label="Ответ", interactive=False)
+    
+    submit_btn.click(
+        fn=get_answer,
+        inputs=question,
+        outputs=answer
+    )
 
-gr.Interface(
-    fn=ask_question,
-    inputs=gr.Textbox(label="Спроси что-нибудь по лору"),
-    outputs=gr.Textbox(label="Ответ"),
-    title="Лор-бот"
-).launch()
+demo.launch(server_name="0.0.0.0", server_port=7860)