Update app.py

app.py

@@ -1,58 +1,82 @@
 import os
+import re
 from langchain_community.document_loaders import TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+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 langchain_community.llms import HuggingFaceHub
+from transformers import pipeline
 import gradio as gr
-import re
 
-# 1. Загрузка и очистка всех .txt файлов
-def load_documents(folder_path):
-    documents = []
-    for file_name in os.listdir(folder_path):
-        if file_name.endswith(".txt"):
-            loader = TextLoader(os.path.join(folder_path, file_name), encoding="utf-8")
-            docs = loader.load()
-            for doc in docs:
-                # Очищаем спецсимволы типа [=/ и прочую ерунду
-                doc.page_content = re.sub(r'\[=/.*?\]', '', doc.page_content)
-                documents.append(doc)
-    return documents
-
-# 2. Разбивка на чанки
-def split_documents(documents):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
-    return splitter.split_documents(documents)
-
-# 3. Создание эмбеддингов
+# 1. Загрузка всех файлов из папки lore/
+def load_all_lore_files():
+    docs = []
+    for filename in os.listdir("lore"):
+        if filename.endswith(".txt"):
+            loader = TextLoader(os.path.join("lore", filename), encoding="utf-8")
+            docs.extend(loader.load())
+    return docs
+
+# 2. Очистка от спецсимволов вроде [=/ и т.п.
+def clean_text(text):
+    return re.sub(r"\[=.*?\/?]", "", text)
+
+# 3. Настройка эмбеддингов
 def create_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
 
-# 4. Загрузка модели
-def load_llm():
-    return HuggingFaceHub(
-        repo_id="IlyaGusev/saiga_mistral_7b_gguf",  # можно заменить на что-то другое, если будет падать
-        model_kwargs={"temperature": 0.6, "max_new_tokens": 300}
-    )
+# 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. Загрузка модели ответа (без HuggingFace API Token)
+def create_llm_pipeline():
+    return pipeline("text-generation", model="IlyaGusev/saiga2_7b_lora", device=0 if torch.cuda.is_available() else -1)
 
-# 5. Построение цепочки
-def build_qa_chain():
-    raw_docs = load_documents("lore")  # Папка lore/ рядом с app.py
-    docs = split_documents(raw_docs)
+# 6. Объединение в цепочку
+def build_chain():
+    docs = load_all_lore_files()
     embeddings = create_embeddings()
-    db = FAISS.from_documents(docs, embeddings)
-    retriever = db.as_retriever()
-    llm = load_llm()
-    return RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
+    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"]
+    )
+
+    return RetrievalQA.from_chain_type(
+        llm=create_llm_pipeline(),
+        retriever=retriever,
+        chain_type_kwargs={"prompt": prompt}
+    )
 
-# 6. Интерфейс
-qa_chain = build_qa_chain()
+# 7. Интерфейс
+qa_chain = build_chain()
 
-def answer_question(question):
-    result = qa_chain.run(question)
-    return result
+def ask_question(question):
+    return qa_chain.run(question)
 
-iface = gr.Interface(fn=answer_question, inputs="text", outputs="text", title="Чат по Лору (RU)")
-iface.launch()
+gr.Interface(
+    fn=ask_question,
+    inputs=gr.Textbox(label="Спроси что-нибудь по лору"),
+    outputs=gr.Textbox(label="Ответ"),
+    title="Лор-бот"
+).launch()