Update app.py

app.py

@@ -1,82 +1,57 @@
 import os
-import re
-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
+import torch
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
 from langchain.chains import RetrievalQA
+from langchain.llms import HuggingFacePipeline
 from transformers import pipeline
-import gradio as gr
 
-# 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
+# Убедись, что все нужные модели и данные сохранены в этой папке
+PERSIST_DIRECTORY = "db"
 
-# 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 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)
-
-# 6. Объединение в цепочку
-def build_chain():
-    docs = load_all_lore_files()
+# Загрузка векторного хранилища
+def load_vectorstore():
     embeddings = create_embeddings()
-    vectorstore = create_vectorstore(docs, embeddings)
-
-    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-
-    prompt = PromptTemplate(
-        template="""
-Ты — помощник, который отвечает на вопросы по вымышленному лору. Отвечай кратко, точно и на русском языке.
-Если в лоре нет нужной информации, честно скажи, что не знаешь.
+    return Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)
 
-Контекст:
-{context}
-
-Вопрос:
-{question}
-
-Ответ:
-""",
-        input_variables=["context", "question"]
+# Создание пайплайна LLM (используем 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
     )
 
+# Обёртка LLM для использования с LangChain
+def create_llm():
+    pipe = create_llm_pipeline()
+    return HuggingFacePipeline(pipeline=pipe)
+
+# Сборка цепочки QA
+def build_chain():
+    vectorstore = load_vectorstore()
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+    llm = create_llm()
     return RetrievalQA.from_chain_type(
-        llm=create_llm_pipeline(),
+        llm=llm,
         retriever=retriever,
-        chain_type_kwargs={"prompt": prompt}
+        return_source_documents=True
     )
 
-# 7. Интерфейс
+# Создаём QA цепочку
 qa_chain = build_chain()
 
-def ask_question(question):
-    return qa_chain.run(question)
-
-gr.Interface(
-    fn=ask_question,
-    inputs=gr.Textbox(label="Спроси что-нибудь по лору"),
-    outputs=gr.Textbox(label="Ответ"),
-    title="Лор-бот"
-).launch()
+# Пример запроса (можешь удалить или адаптировать)
+if __name__ == "__main__":
+    question = "Кто такой Виктор Цой?"
+    result = qa_chain(question)
+    print("\nОтвет:\n", result["result"])