Spaces:
Runtime error
Runtime error
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import gradio as gr
|
3 |
+
from datasets import load_dataset, concatenate_datasets
|
4 |
+
from langchain.docstore.document import Document
|
5 |
+
from langchain.vectorstores import FAISS
|
6 |
+
from langchain_huggingface import HuggingFaceEmbeddings
|
7 |
+
from tqdm import tqdm
|
8 |
+
|
9 |
+
# Путь для сохранения FAISS-индекса
|
10 |
+
INDEX_PATH = "./faiss_index"
|
11 |
+
|
12 |
+
# Инициализируем эмбеддинг-модель (используем модель из Hugging Face)
|
13 |
+
embeddings = HuggingFaceEmbeddings(model_name="fitlemon/bge-m3-uz-legal-matryoshka")
|
14 |
+
|
15 |
+
|
16 |
+
def update_faiss_index():
|
17 |
+
"""
|
18 |
+
Загружает датасет, преобразует данные в документы с метаданными,
|
19 |
+
создаёт FAISS-индекс и сохраняет его локально.
|
20 |
+
"""
|
21 |
+
# Загружаем датасет (например, сплит "train")
|
22 |
+
train_dataset = load_dataset("fitlemon/rag-labor-codex-dataset")["train"]
|
23 |
+
test_dataset = load_dataset("fitlemon/rag-labor-codex-dataset")["test"]
|
24 |
+
# combine train and test datasets
|
25 |
+
dataset = concatenate_datasets([train_dataset, test_dataset])
|
26 |
+
# get rid off duplicate chunks
|
27 |
+
|
28 |
+
docs = []
|
29 |
+
unique_chunks = set()
|
30 |
+
for row in tqdm(dataset, desc="Загрузка документов..."):
|
31 |
+
chunk = row["chunk"]
|
32 |
+
# Если chunk уже добавлен, пропускаем его
|
33 |
+
if chunk in unique_chunks:
|
34 |
+
continue
|
35 |
+
unique_chunks.add(chunk)
|
36 |
+
|
37 |
+
doc = Document(
|
38 |
+
page_content=chunk,
|
39 |
+
metadata={
|
40 |
+
"section": row["section"],
|
41 |
+
"section_name": row["section_name"],
|
42 |
+
"chapter_name": row["chapter"],
|
43 |
+
},
|
44 |
+
)
|
45 |
+
docs.append(doc)
|
46 |
+
|
47 |
+
print(f"Документы успешно загружены и преобразованы. Длина документов: {len(docs)}")
|
48 |
+
# Создаём FAISS-индекс на основе документов
|
49 |
+
db = FAISS.from_documents(docs, embeddings)
|
50 |
+
|
51 |
+
# Сохраняем индекс в указанную директорию
|
52 |
+
os.makedirs(INDEX_PATH, exist_ok=True)
|
53 |
+
db.save_local(INDEX_PATH)
|
54 |
+
print("FAISS индекс обновлён и сохранён в:", INDEX_PATH)
|
55 |
+
return db
|
56 |
+
|
57 |
+
|
58 |
+
# Если индекс ещё не создан, обновляем его, иначе загружаем существующий
|
59 |
+
if not os.path.exists(INDEX_PATH):
|
60 |
+
db = update_faiss_index()
|
61 |
+
else:
|
62 |
+
db = FAISS.load_local(INDEX_PATH, embeddings, allow_dangerous_deserialization=True)
|
63 |
+
print("Загружен существующий FAISS индекс из:", INDEX_PATH)
|
64 |
+
|
65 |
+
|
66 |
+
def retrieve_articles(query):
|
67 |
+
"""
|
68 |
+
Принимает запрос пользователя, ищет в FAISS-индексе топ-3 наиболее релевантных документа
|
69 |
+
и возвращает отформатированный результат в Markdown.
|
70 |
+
"""
|
71 |
+
# Поиск по индексу: возвращает список из документов
|
72 |
+
results = db.similarity_search(query, k=3)
|
73 |
+
|
74 |
+
# Форматируем результаты для вывода
|
75 |
+
result_text = ""
|
76 |
+
for doc in results:
|
77 |
+
result_text += (
|
78 |
+
f"### Статья {doc.metadata['section']}: {doc.metadata['section_name']}\n"
|
79 |
+
)
|
80 |
+
result_text += f"**Глава:** {doc.metadata['chapter_name']}\n\n"
|
81 |
+
result_text += f"**Текст статьи:**\n{doc.page_content}\n\n"
|
82 |
+
result_text += "---\n\n"
|
83 |
+
return result_text
|
84 |
+
|
85 |
+
|
86 |
+
# Создаём Gradio-интерфейс
|
87 |
+
iface = gr.Interface(
|
88 |
+
fn=retrieve_articles,
|
89 |
+
inputs=gr.Textbox(lines=3, placeholder="Введите ваш вопрос о кодексе..."),
|
90 |
+
outputs=gr.Markdown(),
|
91 |
+
title="Поиск по Кодексу через FAISS",
|
92 |
+
description="Введите вопрос, и получите топ-3 наиболее релевантные статьи из кодекса.",
|
93 |
+
)
|
94 |
+
|
95 |
+
if __name__ == "__main__":
|
96 |
+
iface.launch()
|