fitlemon commited on
Commit
0f8f1e7
·
verified ·
1 Parent(s): 79238a7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +96 -0
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()