Update app.py

app.py

@@ -1,12 +1,12 @@
 import os
 import io
-import json
 
 import numpy as np
 import pandas as pd
 import streamlit as st
 from sentence_transformers import SentenceTransformer
-from huggingface_hub import hf_hub_download
+from datasets import load_dataset
+from huggingface_hub import login
 
 # ==========================
 # НАСТРОЙКИ ПРИЛОЖЕНИЯ
@@ -17,43 +17,33 @@ st.set_page_config(
     layout="wide",
 )
 
-# ---- Hugging Face ----
 HF_TOKEN = os.getenv("HF_TOKEN")
-HF_DATA_REPO = os.getenv("HF_DATA_REPO", "yogl/dissertation-search-data")
-
-ALL_DATA_FILENAME = "ALL_data.xlsx"
-EMBEDDINGS_FILENAME = "embeddings_full.json"
+HF_MERGED_REPO = os.getenv("HF_MERGED_REPO", "yogl/dissertations-merged")
+HF_EMB_REPO = os.getenv("HF_EMB_REPO", "yogl/rosrid-dissertations-embeddings")
 
 MODEL_NAME = os.getenv("MODEL_NAME", "intfloat/multilingual-e5-base")
 
-# Порядок колонок в выдаче (внутренние имена)
-COLUMNS_ORDER = [
-    "registration_number",
-    "name",
-    "author_surname",
-    "author_name",
-    "author_patronymic",
-    "author_organization__short_name",
-    "dissertation_type",
-    "created_date",
-    "abstract",
-]
-
 # Отображаемые заголовки (русские имена колонок)
 COLUMN_LABELS_RU = {
     "№": "№",
     "score": "Сходство",
-    "registration_number": "Регистрационный номер",
-    "name": "Название работы",
-    "author_surname": "Фамилия",
-    "author_name": "Имя",
-    "author_patronymic": "Отчество",
-    "author_organization__short_name": "Организация",
+    "fio": "ФИО",
+    "title": "Название диссертации",
+    "author_org_short": "Организация",
     "dissertation_type": "Тип",
-    "created_date": "Дата создания",
-    "abstract": "Аннотация",
+    "registration_number": "Регистрационный номер",
 }
 
+DISPLAY_COLUMNS = [
+    "№",
+    "score",
+    "fio",
+    "title",
+    "author_org_short",
+    "dissertation_type",
+    "registration_number",
+]
+
 if HF_TOKEN is None:
     st.error(
         "Не найден секрет `HF_TOKEN`. "
@@ -61,43 +51,48 @@ if HF_TOKEN is None:
     )
     st.stop()
 
+# Логин к Hugging Face Hub (для приватных датасетов, если нужно)
+try:
+    login(token=HF_TOKEN)
+except Exception:
+    # Если что-то не так с токеном — Streamlit ниже отловит при загрузке датасета
+    pass
+
 # ==========================
 # ЗАГРУЗКА ДАННЫХ И МОДЕЛИ
 # ==========================
 
-@st.cache_data(show_spinner="Скачиваю и загружаю данные из приватного датасета...")
+@st.cache_data(show_spinner="Скачиваю и загружаю данные из Hugging Face датасетов...")
 def load_data():
-    all_data_path = hf_hub_download(
-        repo_id=HF_DATA_REPO,
-        filename=ALL_DATA_FILENAME,
-        repo_type="dataset",
-        token=HF_TOKEN,
+    # 1) МЕТА-ДАННЫЕ: yogl/dissertations-merged
+    ds_meta = load_dataset(
+        HF_MERGED_REPO,
+        split="train",
     )
-
-    embeddings_path = hf_hub_download(
-        repo_id=HF_DATA_REPO,
-        filename=EMBEDDINGS_FILENAME,
-        repo_type="dataset",
-        token=HF_TOKEN,
+    df_meta = ds_meta.to_pandas()
+    # Приводим registration_number к строке и делаем индексом
+    df_meta["registration_number"] = df_meta["registration_number"].astype(str)
+    df_meta = df_meta.set_index("registration_number", drop=False)
+
+    # 2) ЭМБЕДДИНГИ: yogl/rosrid-dissertations-embeddings
+    ds_emb = load_dataset(
+        HF_EMB_REPO,
+        split="train",
     )
+    df_emb = ds_emb.to_pandas()
 
-    df_all = pd.read_excel(all_data_path)
-    df_all["registration_number"] = df_all["registration_number"].astype(str)
-    df_all = df_all.set_index("registration_number", drop=False)
-
-    with open(embeddings_path, "r", encoding="utf-8") as f:
-        emb_data = json.load(f)
+    # Ожидаем поля: registration_number, embedding (список float)
+    df_emb["registration_number"] = df_emb["registration_number"].astype(str)
+    reg_nums = df_emb["registration_number"].tolist()
 
-    reg_nums = [str(item["registration_number"]) for item in emb_data]
-    emb_matrix = np.array(
-        [item["embedding"] for item in emb_data],
-        dtype="float32",
-    )
+    # Превращаем список листов в матрицу
+    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_all, reg_nums, emb_matrix
+    return df_meta, reg_nums, emb_matrix
 
 
 @st.cache_resource(show_spinner="Загружаю модель для эмбеддингов...")
@@ -123,6 +118,7 @@ def search_core(query: str, top_k: int = 10):
     if not query:
         return []
 
+    # e5: запросы с префиксом "query:"
     query_text = "query: " + query
 
     q_emb = model.encode(
@@ -160,46 +156,70 @@ def build_result_df(results):
         else:
             meta = pd.Series({}, index=df_all.columns)
 
+        # Базовые поля
         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),
+            "registration_number": meta.get("registration_number", reg),
+            # Скрытая колонка для построения ссылок
+            "vak_link": meta.get("vak_link", ""),
         }
 
-        for col in COLUMNS_ORDER:
-            row[col] = meta.get(col, None)
-
         rows.append(row)
 
     if not rows:
-        return pd.DataFrame(columns=["№", "score"] + COLUMNS_ORDER)
+        return pd.DataFrame(columns=DISPLAY_COLUMNS + ["vak_link"])
 
     df_res = pd.DataFrame(rows)
-    df_res = df_res[["№", "score"] + COLUMNS_ORDER]
+    # Жёстко задаём порядок колонок
+    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=["№", "score"] + COLUMNS_ORDER)
+        empty_df = pd.DataFrame(columns=DISPLAY_COLUMNS + ["vak_link"])
         empty_df_ru = empty_df.rename(columns=COLUMN_LABELS_RU)
-        df_display = (
-            empty_df_ru.set_index("№") if "№" in empty_df_ru.columns else empty_df_ru
-        )
-        return df_display, None
+        return empty_df_ru, None
 
     results = search_core(query, top_k)
     df_res = build_result_df(results)
 
-    df_res_ru = df_res.rename(columns=COLUMN_LABELS_RU)
-    df_display = df_res_ru.set_index("№")
+    if df_res.empty:
+        empty_df_ru = df_res.rename(columns=COLUMN_LABELS_RU)
+        return empty_df_ru, None
+
+    # --- Разделяем датафреймы для отображения и для Excel ---
+    df_excel = df_res.drop(columns=["vak_link"]).copy()
+    df_display = df_res.copy()
+
+    # В колонке "Тип" делаем HTML-ссылку, если есть vak_link
+    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)
+
+    # Готовим Excel-байты (без HTML-тегов)
     output = io.BytesIO()
     with pd.ExcelWriter(output, engine="xlsxwriter") as writer:
-        df_res_ru.to_excel(writer, index=False)
+        df_excel_ru.to_excel(writer, index=False)
     output.seek(0)
 
-    return df_display, output
+    return df_display_ru, output
 
 
 # ==========================
@@ -240,21 +260,25 @@ with btn_col3:
 
 if clear:
     st.session_state.clear()
-    st.experimental_rerun()
+    try:
+        st.experimental_rerun()
+    except Exception:
+        st.rerun()
 
 # Результаты
 if do_search:
     with st.spinner("Идёт поиск по базе диссертаций..."):
-        df_res, excel_bytes = run_search(query, top_k)
+        df_res_ru, excel_bytes = run_search(query, top_k)
 
-    if df_res.empty:
+    if df_res_ru.empty:
         st.warning("Ничего не найдено. Попробуйте изменить формулировку запроса.")
     else:
-        st.success(f"Найдено записей: {len(df_res)}")
+        st.success(f"Найдено записей: {len(df_res_ru)}")
 
-        st.dataframe(
-            df_res,
-            use_container_width=True,
+        # Показываем таблицу с HTML-ссылками в колонке "Тип"
+        st.markdown(
+            df_res_ru.to_html(escape=False, index=False),
+            unsafe_allow_html=True,
         )
 
         if excel_bytes is not None: