Update app.py

app.py

@@ -1,57 +1,90 @@
 import os
-import torch
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import Chroma
+import re
+import torch  # Добавлен импорт torch
+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 langchain.llms import HuggingFacePipeline
 from transformers import pipeline
+import gradio as gr
 
-# Убедись, что все нужные модели и данные сохранены в этой папке
-PERSIST_DIRECTORY = "db"
+# 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")
+    return HuggingFaceEmbeddings(
+        model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
+    )
 
-# Загрузка векторного хранилища
-def load_vectorstore():
-    embeddings = create_embeddings()
-    return Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)
+# 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)
 
-# Создание пайплайна LLM (используем GPU, если доступен)
+# 5. Загрузка модели ответа (с проверкой доступности GPU)
 def create_llm_pipeline():
     return pipeline(
         "text-generation",
         model="IlyaGusev/saiga2_7b_lora",
-        device=0 if torch.cuda.is_available() else -1,
-        max_new_tokens=512,
-        do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        temperature=0.7
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device="cuda:0" if torch.cuda.is_available() else "cpu"
     )
 
-# Обёртка LLM для использования с LangChain
-def create_llm():
-    pipe = create_llm_pipeline()
-    return HuggingFacePipeline(pipeline=pipe)
-
-# Сборка цепочки QA
+# 6. Объединение в цепочку
 def build_chain():
-    vectorstore = load_vectorstore()
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-    llm = create_llm()
+    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"]
+    )
+
     return RetrievalQA.from_chain_type(
-        llm=llm,
+        llm=create_llm_pipeline(),
         retriever=retriever,
-        return_source_documents=True
+        chain_type_kwargs={"prompt": prompt}
     )
 
-# Создаём QA цепочку
+# 7. Интерфейс
 qa_chain = build_chain()
 
-# Пример запроса (можешь удалить или адаптировать)
-if __name__ == "__main__":
-    question = "Кто такой Виктор Цой?"
-    result = qa_chain(question)
-    print("\nОтвет:\n", result["result"])
+def ask_question(question):
+    return qa_chain.run(question)
+
+gr.Interface(
+    fn=ask_question,
+    inputs=gr.Textbox(label="Спроси что-нибудь по лору"),
+    outputs=gr.Textbox(label="Ответ"),
+    title="Лор-бот"
+).launch()