lepidus / app.py
Loversofdeath's picture
Update app.py
a73e1ef verified
raw
history blame
2.29 kB
import os
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain_community.llms import HuggingFaceHub
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. Создание эмбеддингов
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}
)
# 5. Построение цепочки
def build_qa_chain():
raw_docs = load_documents("lore") # Папка lore/ рядом с app.py
docs = split_documents(raw_docs)
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)
# 6. Интерфейс
qa_chain = build_qa_chain()
def answer_question(question):
result = qa_chain.run(question)
return result
iface = gr.Interface(fn=answer_question, inputs="text", outputs="text", title="Чат по Лору (RU)")
iface.launch()