yogl commited on
Commit
37e80a1
·
verified ·
1 Parent(s): fe3f62f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +180 -92
app.py CHANGED
@@ -1,25 +1,38 @@
1
- import json
2
  import os
3
- import tempfile
 
4
 
5
  import numpy as np
6
  import pandas as pd
7
- import gradio as gr
8
  from sentence_transformers import SentenceTransformer
 
 
 
 
 
9
 
10
- # ==== ФАЙЛЫ ВНУТРИ SPACE ====
11
- ALL_DATA_PATH = "ALL_data.xlsx" # лежит в корне репозитория
12
- EMBEDDINGS_PATH = "embeddings_full.json" # общий файл эмбеддингов
13
- # =============================
14
 
15
- # ==== ЗАГРУЖАЕМ ДАННЫЕ ====
 
 
16
 
17
- print("Загружаю Excel с метаданными...")
18
- df_all = pd.read_excel(ALL_DATA_PATH)
 
19
 
20
- df_all["registration_number"] = df_all["registration_number"].astype(str)
21
- df_all = df_all.set_index("registration_number", drop=False)
 
22
 
 
 
 
 
23
  COLUMNS_ORDER = [
24
  "registration_number",
25
  "name",
@@ -32,22 +45,75 @@ COLUMNS_ORDER = [
32
  "abstract",
33
  ]
34
 
35
- print("Загружаю эмбеддинги...")
36
- with open(EMBEDDINGS_PATH, "r", encoding="utf-8") as f:
37
- emb_data = json.load(f)
38
 
39
- reg_nums = [str(item["registration_number"]) for item in emb_data]
40
- emb_matrix = np.array([item["embedding"] for item in emb_data], dtype="float32")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
 
42
- norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
43
- emb_matrix = emb_matrix / np.maximum(norms, 1e-8)
44
 
45
- MODEL_NAME = "intfloat/multilingual-e5-base"
46
- print(f"Загружаю модель {MODEL_NAME}...")
47
- model = SentenceTransformer(MODEL_NAME)
 
48
 
49
 
50
- # ==== ЛОГИКА ПОИСКА ====
 
 
 
 
 
 
 
 
 
 
51
 
52
  def search_core(query: str, top_k: int = 10):
53
  query = query.strip()
@@ -68,11 +134,13 @@ def search_core(query: str, top_k: int = 10):
68
 
69
  results = []
70
  for rank, idx in enumerate(top_idx, start=1):
71
- results.append({
72
- "rank": rank,
73
- "registration_number": reg_nums[idx],
74
- "score": float(scores[idx]),
75
- })
 
 
76
 
77
  return results
78
 
@@ -95,10 +163,7 @@ def build_result_df(results):
95
  }
96
 
97
  for col in COLUMNS_ORDER:
98
- if col in meta.index:
99
- row[col] = meta[col]
100
- else:
101
- row[col] = None
102
 
103
  rows.append(row)
104
 
@@ -111,11 +176,6 @@ def build_result_df(results):
111
 
112
 
113
  def run_search(query: str, top_k: int):
114
- """
115
- Возвращает:
116
- - DataFrame с результатами
117
- - путь к Excel-файлу с этими результатами
118
- """
119
  query = query.strip()
120
  if not query:
121
  empty_df = pd.DataFrame(columns=["№", "score"] + COLUMNS_ORDER)
@@ -124,66 +184,94 @@ def run_search(query: str, top_k: int):
124
  results = search_core(query, top_k)
125
  df_res = build_result_df(results)
126
 
127
- # Сохраняем во временный Excel
128
- fd, excel_path = tempfile.mkstemp(suffix=".xlsx")
129
- os.close(fd)
130
- df_res.to_excel(excel_path, index=False)
131
-
132
- return df_res, excel_path
133
-
134
 
135
- def clear_all():
136
- """
137
- Очистка: пустой запрос, пустая таблица, отсутствие файла.
138
- """
139
- empty_df = pd.DataFrame(columns=["№", "score"] + COLUMNS_ORDER)
140
- return "", empty_df, None
141
 
142
 
143
- # ==== UI НА BLOCKS ====
 
 
144
 
145
- with gr.Blocks() as demo:
146
- gr.Markdown("## Поиск релевантных диссертаций")
 
 
 
147
 
148
- with gr.Row():
149
- query = gr.Textbox(
150
- label="Запрос",
151
- lines=3,
152
- placeholder="Например: пластификаторы для самоуплотняющихся бетонов",
153
- )
154
-
155
- top_k = gr.Slider(
156
- label="Сколько результатов показать",
157
- minimum=1,
158
- maximum=100,
159
- step=1,
160
- value=20,
161
  )
162
-
163
- with gr.Row():
164
- search_btn = gr.Button("Поиск")
165
- clear_btn = gr.Button("Очистить")
166
-
167
- # ТАБЛИЦА ПОД КНОПКАМИ
168
- results_df = gr.Dataframe(
169
- label="Результаты поиска",
170
- interactive=True, # можно выделять и копировать текст
171
- )
172
- excel_file = gr.File(label="Скачать таблицу (Excel)")
173
-
174
- # Клик по "Поиск"
175
- search_btn.click(
176
- fn=run_search,
177
- inputs=[query, top_k],
178
- outputs=[results_df, excel_file],
179
- )
180
-
181
- # Клик по "Очистить"
182
- clear_btn.click(
183
- fn=clear_all,
184
- inputs=[],
185
- outputs=[query, results_df, excel_file],
186
  )
187
 
188
- if __name__ == "__main__":
189
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import os
2
+ import io
3
+ import json
4
 
5
  import numpy as np
6
  import pandas as pd
7
+ import streamlit as st
8
  from sentence_transformers import SentenceTransformer
9
+ from huggingface_hub import hf_hub_download
10
+
11
+ # ==========================
12
+ # НАСТРОЙКИ ПРИЛОЖЕНИЯ
13
+ # ==========================
14
 
15
+ st.set_page_config(
16
+ page_title="Поиск релевантных диссертаций",
17
+ layout="wide",
18
+ )
19
 
20
+ # ---- Hugging Face ----
21
+ # Секрет HF_TOKEN нужно задать в Settings → Variables and secrets вашего Space.
22
+ HF_TOKEN = os.getenv("HF_TOKEN")
23
 
24
+ # repo_id приватного датасета с файлами ALL_data.xlsx и embeddings_full.json
25
+ # Можешь зашить строкой или задать в переменной окружения HF_DATA_REPO.
26
+ HF_DATA_REPO = os.getenv("HF_DATA_REPO", "username/dissertation-search-data")
27
 
28
+ # Имена файлов в датасете
29
+ ALL_DATA_FILENAME = "ALL_data.xlsx"
30
+ EMBEDDINGS_FILENAME = "embeddings_full.json"
31
 
32
+ # Модель для эмбеддингов
33
+ MODEL_NAME = os.getenv("MODEL_NAME", "intfloat/multilingual-e5-base")
34
+
35
+ # Порядок колонок в выдаче
36
  COLUMNS_ORDER = [
37
  "registration_number",
38
  "name",
 
45
  "abstract",
46
  ]
47
 
 
 
 
48
 
49
+ if HF_TOKEN is None:
50
+ st.error(
51
+ "Не найден секрет `HF_TOKEN`. "
52
+ "Задайте его в Settings → Variables and secrets вашего Space."
53
+ )
54
+ st.stop()
55
+
56
+
57
+ # ==========================
58
+ # ЗАГРУЗКА ДАННЫХ И МОДЕЛИ
59
+ # ==========================
60
+
61
+ @st.cache_data(show_spinner="Скачиваю и загружаю данные из приватного датасета...")
62
+ def load_data():
63
+ # Скачиваем файлы из приватного датасета с помощью токена
64
+ all_data_path = hf_hub_download(
65
+ repo_id=HF_DATA_REPO,
66
+ filename=ALL_DATA_FILENAME,
67
+ repo_type="dataset",
68
+ token=HF_TOKEN,
69
+ )
70
+
71
+ embeddings_path = hf_hub_download(
72
+ repo_id=HF_DATA_REPO,
73
+ filename=EMBEDDINGS_FILENAME,
74
+ repo_type="dataset",
75
+ token=HF_TOKEN,
76
+ )
77
+
78
+ # Загружаем Excel с метаданными
79
+ df_all = pd.read_excel(all_data_path)
80
+ df_all["registration_number"] = df_all["registration_number"].astype(str)
81
+ df_all = df_all.set_index("registration_number", drop=False)
82
+
83
+ # Загружаем эмбеддинги
84
+ with open(embeddings_path, "r", encoding="utf-8") as f:
85
+ emb_data = json.load(f)
86
+
87
+ reg_nums = [str(item["registration_number"]) for item in emb_data]
88
+ emb_matrix = np.array(
89
+ [item["embedding"] for item in emb_data],
90
+ dtype="float32",
91
+ )
92
+
93
+ # Нормируем эмбеддинги
94
+ norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
95
+ emb_matrix = emb_matrix / np.maximum(norms, 1e-8)
96
+
97
+ return df_all, reg_nums, emb_matrix
98
 
 
 
99
 
100
+ @st.cache_resource(show_spinner="Загружаю модель для эмбеддингов...")
101
+ def load_model():
102
+ model = SentenceTransformer(MODEL_NAME)
103
+ return model
104
 
105
 
106
+ try:
107
+ df_all, reg_nums, emb_matrix = load_data()
108
+ model = load_model()
109
+ except Exception as e:
110
+ st.error(f"Ошибка при загрузке данных или модели: {e}")
111
+ st.stop()
112
+
113
+
114
+ # ==========================
115
+ # ЛОГИКА ПОИСКА
116
+ # ==========================
117
 
118
  def search_core(query: str, top_k: int = 10):
119
  query = query.strip()
 
134
 
135
  results = []
136
  for rank, idx in enumerate(top_idx, start=1):
137
+ results.append(
138
+ {
139
+ "rank": rank,
140
+ "registration_number": reg_nums[idx],
141
+ "score": float(scores[idx]),
142
+ }
143
+ )
144
 
145
  return results
146
 
 
163
  }
164
 
165
  for col in COLUMNS_ORDER:
166
+ row[col] = meta.get(col, None)
 
 
 
167
 
168
  rows.append(row)
169
 
 
176
 
177
 
178
  def run_search(query: str, top_k: int):
 
 
 
 
 
179
  query = query.strip()
180
  if not query:
181
  empty_df = pd.DataFrame(columns=["№", "score"] + COLUMNS_ORDER)
 
184
  results = search_core(query, top_k)
185
  df_res = build_result_df(results)
186
 
187
+ # Формируем Excel в памяти
188
+ output = io.BytesIO()
189
+ with pd.ExcelWriter(output, engine="xlsxwriter") as writer:
190
+ df_res.to_excel(writer, index=False)
191
+ output.seek(0)
 
 
192
 
193
+ return df_res, output
 
 
 
 
 
194
 
195
 
196
+ # ==========================
197
+ # UI НА STREAMLIT
198
+ # ==========================
199
 
200
+ st.title("Поиск релевантных диссертаций")
201
+ st.caption(
202
+ "Введите текстовый запрос на естественном языке, например: "
203
+ "«пластификаторы для самоуплотняющихся бетонов»."
204
+ )
205
 
206
+ with st.sidebar:
207
+ st.subheader("Параметры поиска")
208
+ query = st.text_area(
209
+ "Запрос",
210
+ height=120,
211
+ placeholder="Например: пластификаторы для самоуплотняющихся бетонов",
 
 
 
 
 
 
 
212
  )
213
+ top_k = st.slider(
214
+ "Сколько результатов показать",
215
+ min_value=1,
216
+ max_value=100,
217
+ value=20,
218
+ step=1,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
219
  )
220
 
221
+ col_btn1, col_btn2 = st.columns(2)
222
+ with col_btn1:
223
+ do_search = st.button("🔍 Поиск", type="primary")
224
+ with col_btn2:
225
+ clear = st.button("🧹 Очистить")
226
+
227
+ if clear:
228
+ # Полный сброс состояния
229
+ st.session_state.clear()
230
+ st.experimental_rerun()
231
+
232
+ # Основная область
233
+ if do_search:
234
+ with st.spinner("Идёт поиск по базе диссертаций..."):
235
+ df_res, excel_bytes = run_search(query, top_k)
236
+
237
+ if df_res.empty:
238
+ st.warning("Ничего не найдено. Попробуйте изменить формулировку запроса.")
239
+ else:
240
+ st.success(f"Найдено записей: {len(df_res)}")
241
+
242
+ # Можно разделить на вкладки: таблица и детали по столбцам
243
+ tab_table, tab_info = st.tabs(["Таблица результатов", "Описание полей"])
244
+
245
+ with tab_table:
246
+ st.dataframe(
247
+ df_res,
248
+ use_container_width=True,
249
+ )
250
+
251
+ if excel_bytes is not None:
252
+ st.download_button(
253
+ label="💾 Скачать результаты в Excel",
254
+ data=excel_bytes,
255
+ file_name="search_results.xlsx",
256
+ mime=(
257
+ "application/vnd.openxmlformats-officedocument."
258
+ "spreadsheetml.sheet"
259
+ ),
260
+ )
261
+
262
+ with tab_info:
263
+ st.markdown(
264
+ """
265
+ **Пояснения к полям:**
266
+
267
+ - `registration_number` — регистрационный номер диссертации
268
+ - `name` — название работы
269
+ - `author_surname`, `author_name`, `author_patronymic` — ФИО автора
270
+ - `author_organization__short_name` — организация автора
271
+ - `dissertation_type` — тип диссертации
272
+ - `created_date` — дата создания / регистрации
273
+ - `abstract` — аннотация
274
+ """
275
+ )
276
+ else:
277
+ st.info("Введите запрос в левом сайдбаре и нажмите кнопку «Поиск».")