Update app.py

app.py

@@ -22,7 +22,6 @@ 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": "Сходство",
@@ -35,7 +34,6 @@ COLUMN_LABELS_RU = {
     "vak_link": "Ссылка ВАК",  # только в Excel
 }
 
-# Порядок колонок в выдаче (для интерфейса)
 DISPLAY_COLUMNS = [
     "№",
     "score",
@@ -54,7 +52,6 @@ if HF_TOKEN is None:
     )
     st.stop()
 
-# Логин в Hugging Face Hub (для приватных датасетов)
 try:
     login(token=HF_TOKEN)
 except Exception:
@@ -82,10 +79,10 @@ def load_data():
     norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
     emb_matrix = emb_matrix / np.maximum(norms, 1e-8)
 
-    # --- ВЫРАВНИВАЕМ МЕТА-ИНФОРМАЦИЮ ПОД ПОРЯДОК reg_nums ---
+    # --- Выравниваем мета-инфо под порядок reg_nums ---
     meta_aligned = df_meta.reindex(reg_nums)
 
-    # dissertation_type -> маски кандидат/доктор (устойчиво к формулировкам)
+    # dissertation_type -> маски кандидат/доктор
     if "dissertation_type" in meta_aligned.columns:
         type_s = meta_aligned["dissertation_type"].fillna("").astype(str)
     else:
@@ -94,7 +91,7 @@ def load_data():
     is_candidate = type_s.str.contains("кандид", case=False, na=False).to_numpy()
     is_doctor = type_s.str.contains("доктор", case=False, na=False).to_numpy()
 
-    # degree_pursued -> массив строк (для np.isin)
+    # degree_pursued -> массив строк
     if "degree_pursued" in meta_aligned.columns:
         degree_arr = meta_aligned["degree_pursued"].fillna("").astype(str).to_numpy()
     else:
@@ -105,8 +102,7 @@ def load_data():
 
 @st.cache_resource(show_spinner="Загрузка модели...")
 def load_model():
-    model = SentenceTransformer(MODEL_NAME)
-    return model
+    return SentenceTransformer(MODEL_NAME)
 
 
 try:
@@ -120,20 +116,21 @@ except Exception as e:
 # ЛОГИКА ФИЛЬТРОВ И ПОИСКА
 # ==========================
 
-def build_filter_mask(type_selected, degree_selected):
+def build_filter_mask(degree_selected, include_doctors: bool):
+    """
+    По умолчанию: только кандидатские.
+    Если include_doctors=True: кандидатские + докторские.
+    Плюс фильтр по degree_pursued (Науки), если выбран.
+    """
     mask = np.ones(len(reg_nums), dtype=bool)
 
     # Тип диссертации
-    type_selected = set(type_selected or [])
-    if type_selected:
-        type_mask = np.zeros(len(reg_nums), dtype=bool)
-        if "Кандидатские" in type_selected:
-            type_mask |= is_candidate
-        if "Докторские" in type_selected:
-            type_mask |= is_doctor
-        mask &= type_mask
-
-    # Тип наук / степень (degree_pursued)
+    type_mask = is_candidate.copy()
+    if include_doctors:
+        type_mask |= is_doctor
+    mask &= type_mask
+
+    # Науки (degree_pursued)
     if degree_selected:
         mask &= np.isin(degree_arr, degree_selected)
 
@@ -146,11 +143,7 @@ def search_core(query: str, top_k: int = 10, mask=None):
         return []
 
     query_text = "query: " + query
-
-    q_emb = model.encode(
-        query_text,
-        normalize_embeddings=True,
-    )
+    q_emb = model.encode(query_text, normalize_embeddings=True)
 
     if mask is None:
         idx_pool = np.arange(len(reg_nums))
@@ -169,17 +162,12 @@ def search_core(query: str, top_k: int = 10, mask=None):
     results = []
     for rank, (idx, sc) in enumerate(zip(top_idx, top_scores), start=1):
         results.append(
-            {
-                "rank": rank,
-                "registration_number": reg_nums[idx],
-                "score": float(sc),
-            }
+            {"rank": rank, "registration_number": reg_nums[idx], "score": float(sc)}
         )
     return results
 
 
 def extract_year(value):
-    """Аккуратно вытаскиваем год защиты из поля protection_date."""
     if value is None:
         return None
     try:
@@ -206,14 +194,12 @@ def build_result_df(results):
 
         if reg in df_all.index:
             meta = df_all.loc[reg]
-            # если вдруг дубликаты по индексу (редко), берём первую строку
-            if isinstance(meta, pd.DataFrame):
+            if isinstance(meta, pd.DataFrame):  # на случай дублей
                 meta = meta.iloc[0]
         else:
             meta = pd.Series({}, index=df_all.columns)
 
-        protection_date = meta.get("protection_date", None)
-        protection_year = extract_year(protection_date)
+        protection_year = extract_year(meta.get("protection_date", None))
 
         row = {
             "№": r["rank"],
@@ -226,38 +212,33 @@ def build_result_df(results):
             "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
+    return df_res[DISPLAY_COLUMNS + ["vak_link"]]
 
 
-def run_search(query: str, top_k: int, type_selected, degree_selected):
+def run_search(query: str, top_k: int, degree_selected, include_doctors: bool):
     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
+        return empty_df.rename(columns=COLUMN_LABELS_RU), None
 
-    mask = build_filter_mask(type_selected, degree_selected)
+    mask = build_filter_mask(degree_selected, include_doctors)
 
     results = search_core(query, top_k, mask=mask)
     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
+        return df_res.rename(columns=COLUMN_LABELS_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 ""
@@ -271,7 +252,6 @@ def run_search(query: str, top_k: int, type_selected, degree_selected):
     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)
@@ -283,13 +263,48 @@ def run_search(query: str, top_k: int, type_selected, degree_selected):
 # UI НА STREAMLIT
 # ==========================
 
-# Заголовок по центру
 st.markdown(
     "<h1 style='text-align: center; margin-bottom: 0.5rem;'>Поиск постдока🎓</h1>",
     unsafe_allow_html=True,
 )
 
-# Форма: Ctrl+Enter = отправка
+def _pick_default_sciences(options):
+    """
+    Требуемые по умолчанию: технические, физмат, химические, биологические.
+    Подбираем по подстрокам (на случай отличий в формулировках).
+    """
+    want = [
+        ("технич",),
+        ("физ", "мат"),          # физ-мат/физмат/физико-математические
+        ("хим",),
+        ("биолог",),
+    ]
+    opts_l = [o.lower() for o in options]
+
+    picked = []
+    for keys in want:
+        found = None
+        for o, ol in zip(options, opts_l):
+            if all(k in ol for k in keys):
+                found = o
+                break
+        if found:
+            picked.append(found)
+
+    # если по каким-то причинам не нашли ничего — оставим первые 4 (чтобы фильтр не был пустым)
+    if not picked and options:
+        picked = options[: min(4, len(options))]
+
+    # убираем дубли, сохраняя порядок
+    seen = set()
+    uniq = []
+    for x in picked:
+        if x not in seen:
+            seen.add(x)
+            uniq.append(x)
+    return uniq
+
+
 with st.form("search_form"):
     top_k = st.slider(
         "Сколько результатов показать",
@@ -306,14 +321,8 @@ with st.form("search_form"):
         key="query",
     )
 
-    # Скрываемые точные настройки
     with st.expander("Точные настройки", expanded=False):
-        type_selected = st.multiselect(
-            "Тип диссертаций для поиска",
-            options=["Кандидатские", "Докторские"],
-            default=["Кандидатские", "Докторские"],
-        )
-
+        # Науки (degree_pursued)
         if "degree_pursued" in df_all.columns:
             degree_options = (
                 df_all["degree_pursued"]
@@ -328,13 +337,20 @@ with st.form("search_form"):
         else:
             degree_options = []
 
+        default_sciences = _pick_default_sciences(degree_options)
+
         degree_selected = st.multiselect(
-            "Тип наук (degree_pursued)",
+            "Науки",
             options=degree_options,
-            default=[],
+            default=default_sciences,
+        )
+
+        # По умолчанию: только кандидатские (чекбокс выключен)
+        include_doctors = st.checkbox(
+            "Включать докторские диссертации",
+            value=False,
         )
 
-    # Кнопка по центру через три колонки
     c1, c2, c3 = st.columns([1, 1, 1])
     with c1:
         st.write("")
@@ -343,10 +359,9 @@ with st.form("search_form"):
     with c3:
         st.write("")
 
-# Результаты
 if do_search:
     with st.spinner("Идёт поиск по базе диссертаций..."):
-        df_res_ru, excel_bytes = run_search(query, top_k, type_selected, degree_selected)
+        df_res_ru, excel_bytes = run_search(query, top_k, degree_selected, include_doctors)
 
     if df_res_ru.empty:
         st.warning("Ничего не найдено. Проверьте запрос и/или ослабьте фильтры в «Точных настройках».")
@@ -382,10 +397,7 @@ if do_search:
                 label="💾 Скачать результаты в Excel",
                 data=excel_bytes,
                 file_name="search_results.xlsx",
-                mime=(
-                    "application/vnd.openxmlformats-officedocument."
-                    "spreadsheetml.sheet"
-                ),
+                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
             )
 else:
     st.info(
@@ -393,7 +405,6 @@ else:
         "или нажмите Ctrl+Enter в поле ввода."
     )
 
-# Футер
 st.markdown(
     "<p style='font-size: 0.8rem; text-align: right; color: gray;'>"
     "(с) Антон Лощилов, 2025"