import os import io import numpy as np import pandas as pd import streamlit as st from sentence_transformers import SentenceTransformer from datasets import load_dataset from huggingface_hub import login # ========================== # НАСТРОЙКИ ПРИЛОЖЕНИЯ # ========================== st.set_page_config( page_title="Поиск постдока", layout="wide", ) HF_TOKEN = os.getenv("HF_TOKEN") HF_MERGED_REPO = os.getenv("HF_MERGED_REPO") HF_EMB_REPO = os.getenv("HF_EMB_REPO") MODEL_NAME = os.getenv("MODEL_NAME") # Отображаемые заголовки (русские имена колонок) COLUMN_LABELS_RU = { "№": "№", "score": "Сходство", "fio": "ФИО", "title": "Название диссертации", "author_org_short": "Организация", "dissertation_type": "Тип", "protection_year": "Год", "registration_number": "Регистрационный номер", "vak_link": "Ссылка ВАК", # только в Excel } # Порядок колонок в выдаче (для интерфейса) DISPLAY_COLUMNS = [ "№", "score", "fio", "title", "author_org_short", "dissertation_type", "protection_year", "registration_number", ] if HF_TOKEN is None: st.error( "Не найден секрет `HF_TOKEN`. " "Задайте его в Settings → Variables and secrets вашего Space." ) st.stop() # Логин в Hugging Face Hub (для приватных датасетов) try: login(token=HF_TOKEN) except Exception: pass # ========================== # ЗАГРУЗКА ДАННЫХ И МОДЕЛИ # ========================== @st.cache_data(show_spinner="Загрузка данных...") def load_data(): # 1) МЕТА-ДАННЫЕ ds_meta = load_dataset( HF_MERGED_REPO, split="train", ) df_meta = ds_meta.to_pandas() df_meta["registration_number"] = df_meta["registration_number"].astype(str) df_meta = df_meta.set_index("registration_number", drop=False) # 2) ЭМБЕДДИНГИ ds_emb = load_dataset( HF_EMB_REPO, split="train", ) df_emb = ds_emb.to_pandas() df_emb["registration_number"] = df_emb["registration_number"].astype(str) reg_nums = df_emb["registration_number"].tolist() emb_matrix = np.vstack(df_emb["embedding"].to_list()).astype("float32") norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True) emb_matrix = emb_matrix / np.maximum(norms, 1e-8) return df_meta, reg_nums, emb_matrix @st.cache_resource(show_spinner="Загрузка модели...") def load_model(): model = SentenceTransformer(MODEL_NAME) return model try: df_all, reg_nums, emb_matrix = load_data() model = load_model() except Exception as e: st.error(f"Ошибка при загрузке данных или модели: {e}") st.stop() # ========================== # ЛОГИКА ПОИСКА # ========================== def search_core(query: str, top_k: int = 10): query = query.strip() if not query: return [] query_text = "query: " + query q_emb = model.encode( query_text, normalize_embeddings=True, ) scores = emb_matrix @ q_emb top_k = min(int(top_k), len(scores)) top_idx = np.argsort(-scores)[:top_k] results = [] for rank, idx in enumerate(top_idx, start=1): results.append( { "rank": rank, "registration_number": reg_nums[idx], "score": float(scores[idx]), } ) return results def extract_year(value): """Аккуратно вытаскиваем год защиты из поля protection_date.""" if value is None: return None try: if not isinstance(value, str): if pd.isna(value): return None dt = pd.to_datetime(value) return str(dt.year) except Exception: pass s = str(value).strip() if len(s) >= 4 and s[:4].isdigit(): return s[:4] return None def build_result_df(results): rows = [] for r in results: reg = r["registration_number"] score = r["score"] if reg in df_all.index: meta = df_all.loc[reg] else: meta = pd.Series({}, index=df_all.columns) protection_date = meta.get("protection_date", None) protection_year = extract_year(protection_date) row = { "№": r["rank"], "score": round(score, 4), "fio": meta.get("fio", None), "title": meta.get("title", None), "author_org_short": meta.get("author_org_short", None), "dissertation_type": meta.get("dissertation_type", None), "protection_year": protection_year, "registration_number": meta.get("registration_number", reg), "vak_link": meta.get("vak_link", ""), } rows.append(row) if not rows: return pd.DataFrame(columns=DISPLAY_COLUMNS + ["vak_link"]) df_res = pd.DataFrame(rows) df_res = df_res[DISPLAY_COLUMNS + ["vak_link"]] return df_res def run_search(query: str, top_k: int): query = query.strip() if not query: empty_df = pd.DataFrame(columns=DISPLAY_COLUMNS + ["vak_link"]) empty_df_ru = empty_df.rename(columns=COLUMN_LABELS_RU) return empty_df_ru, None results = search_core(query, top_k) df_res = build_result_df(results) if df_res.empty: empty_df_ru = df_res.rename(columns=COLUMN_LABELS_RU) return empty_df_ru, None # --- Разделяем датафреймы для отображения и для Excel --- df_excel = df_res.copy() # vak_link оставляем df_display = df_res.copy() # HTML-ссылка в колонке "Тип" def make_type_cell(row): t = row.get("dissertation_type", "") link = row.get("vak_link") or "" if isinstance(t, str) and t and link: return f'{t}' return t df_display["dissertation_type"] = df_display.apply(make_type_cell, axis=1) df_display = df_display.drop(columns=["vak_link"]) df_display_ru = df_display.rename(columns=COLUMN_LABELS_RU) df_excel_ru = df_excel.rename(columns=COLUMN_LABELS_RU) # Excel (с отдельной колонкой "Ссылка ВАК") output = io.BytesIO() with pd.ExcelWriter(output, engine="xlsxwriter") as writer: df_excel_ru.to_excel(writer, index=False) output.seek(0) return df_display_ru, output # ========================== # UI НА STREAMLIT # ========================== # Заголовок по центру st.markdown( "

Поиск постдока🎓

", unsafe_allow_html=True, ) # Форма: Ctrl+Enter = отправка with st.form("search_form"): top_k = st.slider( "Сколько результатов показать", min_value=1, max_value=100, value=20, step=1, ) query = st.text_area( "Введите запрос", height=120, placeholder="Например: пластификаторы для самоуплотняющихся бетонов", key="query", ) # Кнопка по центру через три колонки c1, c2, c3 = st.columns([1, 1, 1]) with c1: st.write("") with c2: do_search = st.form_submit_button("🔍 Поиск", type="primary",use_container_width=True ) with c3: st.write("") # Результаты if do_search: with st.spinner("Идёт поиск по базе диссертаций..."): df_res_ru, excel_bytes = run_search(query, top_k) if df_res_ru.empty: st.warning("Ничего не найдено. Попробуйте изменить формулировку запроса.") else: st.success(f"Найдено записей: {len(df_res_ru)}") table_html = df_res_ru.to_html( escape=False, index=False, classes="result-table", ) st.markdown( """ """, unsafe_allow_html=True, ) st.markdown(table_html, unsafe_allow_html=True) if excel_bytes is not None: st.download_button( label="💾 Скачать результаты в Excel", data=excel_bytes, file_name="search_results.xlsx", mime=( "application/vnd.openxmlformats-officedocument." "spreadsheetml.sheet" ), ) else: st.info("Введите запрос выше и нажмите кнопку «Поиск» " "или нажмите Ctrl+Enter в поле ввода.") # Футер st.markdown( "

" "(с) Антон Лощилов, 2025" "

", unsafe_allow_html=True, )