| 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": "Ссылка ВАК", |
| } |
|
|
| |
| 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() |
|
|
| |
| try: |
| login(token=HF_TOKEN) |
| except Exception: |
| pass |
|
|
|
|
| |
| |
| |
|
|
| @st.cache_data(show_spinner="Загрузка данных...") |
| def load_data(): |
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| df_excel = df_res.copy() |
| df_display = df_res.copy() |
|
|
| |
| 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'<a href="{link}" target="_blank">{t}</a>' |
| 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) |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| |
| st.markdown( |
| "<h1 style='text-align: center; margin-bottom: 0.5rem;'>Поиск постдока🎓</h1>", |
| unsafe_allow_html=True, |
| ) |
|
|
| |
| 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( |
| """ |
| <style> |
| table.result-table { |
| width: 100%; |
| border-collapse: collapse; |
| } |
| table.result-table th { |
| text-align: center !important; |
| vertical-align: middle; |
| } |
| </style> |
| """, |
| 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( |
| "<p style='font-size: 0.8rem; text-align: right; color: gray;'>" |
| "(с) Антон Лощилов, 2025" |
| "</p>", |
| unsafe_allow_html=True, |
| ) |
|
|
